gpu-cpu

models_725/segWaterBuilding.py

@@ -16,8 +16,9 @@ class SegModel(object):
                                 use_boundary_2=False, use_boundary_4=False,
                                 use_boundary_8=True, use_boundary_16=False,
                                 use_conv_last=False)
-        self.model.load_state_dict(torch.load(respth))
-        self.device = device
+        self.device = device      
+        self.model.load_state_dict(torch.load(respth, map_location=torch.device(self.device)  ))
+        
         self.multiOutput = multiOutput
         self.model= self.model.to(self.device)
         self.to_tensor = transforms.Compose([
@@ -30,8 +31,9 @@ class SegModel(object):
     def eval(self, image=None):
         H, W, _ = image.shape
         img = self.preprocess_image(image)
-
-        imgs = img.cuda()
+        if self.device != 'cpu':
+            imgs = img.cuda()
+        else:imgs=img    
         self.model.eval()
         logits = self.model(imgs)
 

predict.py

@@ -23,7 +23,8 @@ if __name__ == '__main__':
     impth = 'images/examples'
     outpth= 'images/results'
     folders = os.listdir(impth)
-    segmodel = SegModel()
+    #segmodel = SegModel(device='cuda:0')
+    segmodel = SegModel(device='cpu')
     for i in range(len(folders)):
         imgpath = os.path.join(impth, folders[i])
         time00 = time.time()