enhance(example): update example code with device converter (#1264)

update example code with device converter
star-whale · Sep 21, 2022 · c37fadd · c37fadd
1 parent 8b6e394
commit c37fadd
Show file tree

Hide file tree

Showing 8 changed files with 9 additions and 57 deletions.
diff --git a/example/PennFudanPed/pfp/ppl.py b/example/PennFudanPed/pfp/ppl.py
@@ -36,7 +36,7 @@ def cmp(self, ppl_result):
 
         evaluator = make_coco_evaluator(annotations, iou_types=["bbox", "segm"])
         for index, pred in pred_results:
-            evaluator.update({index: pred})
+            evaluator.update({index: {k: v.cpu() for k, v in pred.items()}})
 
         evaluator.synchronize_between_processes()
         evaluator.accumulate()

diff --git a/example/cifar10/cifar/ppl.py b/example/cifar10/cifar/ppl.py
@@ -59,7 +59,7 @@ def _post(self, input):
 
     def _load_model(self, device):
         model = Net().to(device)
-        model.load_state_dict(torch.load(str(ROOTDIR / "models" / "cifar_net.pth")))
+        model.load_state_dict(torch.load(str(ROOTDIR / "models" / "cifar_net.pth"), map_location=device))
         model.eval()
         print("load cifar_net model, start to inference...")
         return model
diff --git a/example/mnist/mnist/ppl.py b/example/mnist/mnist/ppl.py
@@ -61,7 +61,7 @@ def _post(self, input):
 
     def _load_model(self, device):
         model = Net().to(device)
-        model.load_state_dict(torch.load(str(ROOTDIR / "models/mnist_cnn.pt")))
+        model.load_state_dict(torch.load(str(ROOTDIR / "models/mnist_cnn.pt"), map_location=device))
         model.eval()
         print("load mnist model, start to inference...")
         return model

diff --git a/example/mnist/mnist/test.py b/example/mnist/mnist/test.py
diff --git a/example/nmt/nmt/ppl.py b/example/nmt/nmt/ppl.py
@@ -78,12 +78,12 @@ def _load_vocab(self):
     def _load_encoder_model(self, device):
         model = EncoderRNN(self.vocab.vin.n_words, self.hidden_size, device).to(device)
 
-        param = torch.load(_ROOT_DIR + "/models/encoder.pth", device)
+        param = torch.load(_ROOT_DIR + "/models/encoder.pth", map_location=device)
         model.load_state_dict(param)
         return model
 
     def _load_decoder_model(self, device):
         model = DecoderRNN(self.vocab.vout.n_words, self.hidden_size, device).to(device)
-        param = torch.load(_ROOT_DIR + "/models/decoder.pth", device)
+        param = torch.load(_ROOT_DIR + "/models/decoder.pth", map_location=device)
         model.load_state_dict(param)
         return model
diff --git a/example/runtime/pytorch/runtime.yaml b/example/runtime/pytorch/runtime.yaml
@@ -6,11 +6,11 @@ configs:
   docker:
     image: ghcr.io/star-whale/runtime/pytorch
   pip:
-    extra_index_url: https://pypi.doubanio.com/simple
+    extra_index_url: https://mirrors.bfsu.edu.cn/pypi/web/simple
     index_url: https://pypi.tuna.tsinghua.edu.cn/simple
     trusted_host:
       - pypi.tuna.tsinghua.edu.cn
-      - pypi.doubanio.com
+      - mirrors.bfsu.edu.cn
 dependencies:
   - pip:
       - Pillow

diff --git a/example/speech_command/sc/ppl.py b/example/speech_command/sc/ppl.py
@@ -104,7 +104,7 @@ def _post(self, input: torch.Tensor) -> t.Tuple[t.List[int], t.List[float]]:
 
     def _load_model(self, device):
         model = M5(n_input=1, n_output=len(ALL_LABELS))
-        model.load_state_dict(torch.load(str(ROOTDIR / "models/m5.pth")))
+        model.load_state_dict(torch.load(str(ROOTDIR / "models/m5.pth"), map_location=device))
         model.to(device)
         model.eval()
         print("m5 model loaded, start to inference...")

diff --git a/example/text_cls_AG_NEWS/tcan/ppl.py b/example/text_cls_AG_NEWS/tcan/ppl.py
@@ -52,7 +52,7 @@ def cmp(self, ppl_result):
     def _load_model(self, device):
         model_path = _ROOT_DIR + "/models/model.i"
         model = TextClassificationModel(1308713, 32, _NUM_CLASSES).to(device)
-        model.load_state_dict(torch.load(model_path))
+        model.load_state_dict(torch.load(model_path, map_location=device))
         model.eval()
         return model