CPU Version: fix torch.load

joints_detectors/Alphapose/SPPE/src/main_fast_inference.py

@@ -27,7 +27,7 @@ class InferenNet(nn.Module):
         model = createModel()
         print('Loading pose model from {}'.format('joints_detectors/Alphapose/models/sppe/duc_se.pth'))
         sys.stdout.flush()
-        model.load_state_dict(torch.load('joints_detectors/Alphapose/models/sppe/duc_se.pth'))
+        model.load_state_dict(torch.load('joints_detectors/Alphapose/models/sppe/duc_se.pth', map_location=torch.device('cpu')))
         model.eval()
         self.pyranet = model
 
@@ -54,7 +54,7 @@ class InferenNet_fast(nn.Module):
 
         model = createModel()
         print('Loading pose model from {}'.format('models/sppe/duc_se.pth'))
-        model.load_state_dict(torch.load('models/sppe/duc_se.pth'))
+        model.load_state_dict(torch.load('models/sppe/duc_se.pth', map_location=torch.device('cpu')))
         model.eval()
         self.pyranet = model
 

joints_detectors/Alphapose/SPPE/src/opt.py

@@ -96,7 +96,7 @@ parser.add_argument('--port', dest='port',
 
 opt = parser.parse_args()
 if opt.Continue:
-    opt = torch.load("../exp/{}/{}/option.pkl".format(opt.dataset, opt.expID))
+    opt = torch.load("../exp/{}/{}/option.pkl".format(opt.dataset, opt.expID), map_location=torch.device('cpu'))
     opt.Continue = True
     opt.nEpochs = 50
     print("--- Continue ---")

joints_detectors/Alphapose/train_sppe/src/train.py

@@ -112,7 +112,7 @@ def main():
     m = createModel()
     if opt.loadModel:
         print('Loading Model from {}'.format(opt.loadModel))
-        m.load_state_dict(torch.load(opt.loadModel))
+        m.load_state_dict(torch.load(opt.loadModel, map_location=torch.device('cpu')))
         if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
             try:
                 os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))

tools/utils.py

@@ -138,7 +138,7 @@ def videopose_model_load():
     # load trained model
     from common.model import TemporalModel
     chk_filename = main_path + '/checkpoint/pretrained_h36m_detectron_coco.bin'
-    checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage)  # 把loc映射到storage
+    checkpoint = torch.load(chk_filename, map_location=torch.device('cpu'))  # 把loc映射到storage
     model_pos = TemporalModel(17, 2, 17, filter_widths=[3, 3, 3, 3, 3], causal=False, dropout=False, channels=1024, dense=False)
     model_pos = model_pos
     model_pos.load_state_dict(checkpoint['model_pos'])

videopose_PSTMO.py

@@ -103,7 +103,7 @@ def main(args):
     model_dict = model['trans'].state_dict()
 
     no_refine_path = "checkpoint/PSTMOS_no_refine_48_5137_in_the_wild.pth"
-    pre_dict = torch.load(no_refine_path)
+    pre_dict = torch.load(no_refine_path, map_location=torch.device('cpu'))
     for key, value in pre_dict.items():
         name = key[7:]
         model_dict[name] = pre_dict[key]