Support running on CPU (#25)

Thanks for your Contribution!

benchmark.py

@@ -82,20 +82,25 @@ def do_crf_inference(image, unary, args):
     img = image.transpose(2, 0, 1)  # shape: [3, hight, width]
     # Add batch dimension to image: [1, 3, height, width]
     img = img.reshape([1, 3, shape[0], shape[1]])
-    img_var = Variable(torch.Tensor(img)).cuda()
+    img_var = Variable(torch.Tensor(img))
 
     un = unary.transpose(2, 0, 1)  # shape: [3, hight, width]
     # Add batch dimension to unary: [1, 21, height, width]
     un = un.reshape([1, num_classes, shape[0], shape[1]])
-    unary_var = Variable(torch.Tensor(un)).cuda()
+    unary_var = Variable(torch.Tensor(un))
 
     logging.debug("Build ConvCRF.")
     ##
     # Create CRF module
-    gausscrf = convcrf.GaussCRF(conf=config, shape=shape, nclasses=num_classes)
-    # Cuda computation is required.
-    # A CPU implementation of our message passing is not provided.
-    gausscrf.cuda()
+    gausscrf = convcrf.GaussCRF(conf=config, shape=shape, nclasses=num_classes,
+                                use_gpu=not args.cpu)
+
+    # move to GPU if requested
+    if not args.cpu:
+        img_var = img_var.cuda()
+        unary_var = unary_var.cuda()
+        gausscrf.cuda()
+
 
     # Perform ConvCRF inference
     """
@@ -262,6 +267,9 @@ def get_parser():
                         help="Use pyinn based Cuda implementation"
                              "for message passing.")
 
+    parser.add_argument('--cpu', action='store_true',
+                        help="Run on CPU instead of GPU.")
+
     return parser
 
 

convcrf/convcrf.py

@@ -119,7 +119,7 @@ class GaussCRF(nn.Module):
         "Connected CRFs with Gaussian Edge Pots" (arxiv.org/abs/1210.5644)
     """
 
-    def __init__(self, conf, shape, nclasses=None):
+    def __init__(self, conf, shape, nclasses=None, use_gpu=True):
         super(GaussCRF, self).__init__()
 
         self.conf = conf
@@ -163,7 +163,7 @@ def register(name, tensor):
 
         self.CRF = ConvCRF(
             shape, nclasses, mode="col", conf=conf,
-            use_gpu=True, filter_size=conf['filter_size'],
+            use_gpu=use_gpu, filter_size=conf['filter_size'],
             norm=conf['norm'], blur=conf['blur'], trainable=conf['trainable'],
             convcomp=conf['convcomp'], weight=weight,
             final_softmax=conf['final_softmax'],
@@ -281,8 +281,6 @@ def __init__(self, feat_list, compat_list, merge, npixels, nclasses,
                  filter_size=5, clip_edges=0, use_gpu=False,
                  blur=1, matmul=False, verbose=False, pyinn=False):
 
-        assert(use_gpu)
-
         if not norm == "sym" and not norm == "none":
             raise NotImplementedError
 
@@ -558,16 +556,14 @@ def clean_filters(self):
     def add_pairwise_energies(self, feat_list, compat_list, merge):
         assert(len(feat_list) == len(compat_list))
 
-        assert(self.use_gpu)
-
         self.kernel = MessagePassingCol(
             feat_list=feat_list,
             compat_list=compat_list,
             merge=merge,
             npixels=self.npixels,
             filter_size=self.filter_size,
             nclasses=self.nclasses,
-            use_gpu=True,
+            use_gpu=self.use_gpu,
             norm=self.norm,
             verbose=self.verbose,
             blur=self.blur,

demo.py

@@ -80,20 +80,24 @@ def do_crf_inference(image, unary, args):
     image = image.transpose(2, 0, 1)  # shape: [3, hight, width]
     # Add batch dimension to image: [1, 3, height, width]
     image = image.reshape([1, 3, shape[0], shape[1]])
-    img_var = Variable(torch.Tensor(image)).cuda()
+    img_var = Variable(torch.Tensor(image))
 
     unary = unary.transpose(2, 0, 1)  # shape: [3, hight, width]
     # Add batch dimension to unary: [1, 21, height, width]
     unary = unary.reshape([1, num_classes, shape[0], shape[1]])
-    unary_var = Variable(torch.Tensor(unary)).cuda()
+    unary_var = Variable(torch.Tensor(unary))
 
     logging.info("Build ConvCRF.")
     ##
     # Create CRF module
-    gausscrf = convcrf.GaussCRF(conf=config, shape=shape, nclasses=num_classes)
-    # Cuda computation is required.
-    # A CPU implementation of our message passing is not provided.
-    gausscrf.cuda()
+    gausscrf = convcrf.GaussCRF(conf=config, shape=shape, nclasses=num_classes,
+                                use_gpu=not args.cpu)
+
+    # move to GPU if requested
+    if not args.cpu:
+        img_var = img_var.cuda()
+        unary_var = unary_var.cuda()
+        gausscrf.cuda()
 
     logging.info("Start Computation.")
     # Perform CRF inference
@@ -212,6 +216,9 @@ def get_parser():
                         help="Use pyinn based Cuda implementation"
                              "for message passing.")
 
+    parser.add_argument('--cpu', action='store_true',
+                        help="Run on CPU instead of GPU.")
+
     # parser.add_argument('--compare', action='store_true')
     # parser.add_argument('--embed', action='store_true')