fix multi-gpu bug in lib/non_local.py

Author: AlexHex7

README.md

@@ -4,6 +4,7 @@
 ## Statement
 - Only do the experiments on MNIST dataset so far.
 - You can find the non-local block in **lib/**. 
+- The code can support **multi-gpu** now.
 - If there is something wrong in my code, please contact me, thanks!
 
 There are two version **non-local.py** and **non-local-simple-version.py**. 
@@ -20,3 +21,5 @@ There are two version **non-local.py** and **non-local-simple-version.py**.
 - Experiments on Charades dataset.
 - Experiments on COCO dataset.
 - [x] Make sure how to do the Implementation of concatenation.
+- [x] Support multi-gpu.
+- [x] Fix the bug in **lib/non_local.py** when using multi-gpu.

demo_MNIST.py

@@ -32,6 +32,7 @@ def calc_acc(x, y):
 
 net = Network()
 if torch.cuda.is_available():
+    net = nn.DataParallel(net)
     net.cuda()
 
 opt = torch.optim.Adam(net.parameters(), lr=cfg.LR, weight_decay=cfg.weight_decay)

lib/non_local.py

@@ -59,24 +59,36 @@ def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded
         self.phi = None
         self.concat_project = None
 
-        if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:
-            self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
-                                 kernel_size=1, stride=1, padding=0)
-            self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
-                               kernel_size=1, stride=1, padding=0)
-
-            if mode == 'embedded_gaussian':
-                self.operation_function = self._embedded_gaussian
-            elif mode == 'dot_product':
-                self.operation_function = self._dot_product
-            elif mode == 'concatenation':
-                self.operation_function = self._concatenation
-                self.concat_project = nn.Sequential(
-                    nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
-                    nn.ReLU()
-                )
-        elif mode == 'gaussian':
-            self.operation_function = self._gaussian
+        # if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:
+        self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                             kernel_size=1, stride=1, padding=0)
+
+        self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+                           kernel_size=1, stride=1, padding=0)
+        #       elif mode == 'concatenation':
+        self.concat_project = nn.Sequential(
+            nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
+            nn.ReLU()
+        )
+
+        # if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:
+        #     self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+        #                          kernel_size=1, stride=1, padding=0)
+        #     self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
+        #                        kernel_size=1, stride=1, padding=0)
+        #
+        #     if mode == 'embedded_gaussian':
+        #         self.operation_function = self._embedded_gaussian
+        #     elif mode == 'dot_product':
+        #         self.operation_function = self._dot_product
+        #     elif mode == 'concatenation':
+        #         self.operation_function = self._concatenation
+        #         self.concat_project = nn.Sequential(
+        #             nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
+        #             nn.ReLU()
+        #         )
+        # elif mode == 'gaussian':
+        #     self.operation_function = self._gaussian
 
         if sub_sample:
             self.g = nn.Sequential(self.g, max_pool(kernel_size=2))
@@ -91,7 +103,15 @@ def forward(self, x):
         :return:
         '''
 
-        output = self.operation_function(x)
+        if self.mode == 'embedded_gaussian':
+            output = self._embedded_gaussian(x)
+        elif mode == 'dot_product':
+            output = self._dot_product(x)
+        elif mode == 'concatenation':
+            output = self._concatenation(x)
+        elif mode == 'gaussian':
+            output = self._gaussian(x)
+
         return output
 
     def _embedded_gaussian(self, x):