fix 3dpooling bug and simplify code

Author: AlexHex7

README.md

@@ -1,25 +1,40 @@
 # non-local_pytorch
-- Implementation of [**Non-local Neural Networks**](https://arxiv.org/abs/1711.07971).
+- Implementation of [**Non-local Neural Block**](https://arxiv.org/abs/1711.07971).
 
 ## 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.
+- You can find different kinds of non-local block in **lib/**. 
+- The code is tested on MNIST dataset. You can select the type of non-local block 
+in **lib/network.py**.
 - 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**. 
-
-- **non-local.py** contains the implementation of Gaussian, embedded Gaussian and  dot product, which is mainly for learning.
-- **non-local-simple-version.py** only contains  the implementation of embedded Gaussian.
 
 ## Environment
 - python 3.6
 - pytorch 0.3.0
 
+## Update Records
+1. Figure out how to implement the **concatenation** type, and add the code to **lib/**.
+2. Fix the bug in **lib/non_local.py** (old version) when using multi-gpu. Someone shares the 
+reason with me, and you can find it in [here](https://github.com/pytorch/pytorch/issues/8637).
+3. Fix the bug of 3D pooling in **lib/non_local.py** (old version). Appreciate 
+[**protein27**](https://github.com/AlexHex7/Non-local_pytorch/issues/17) for pointing it out.
+4. For convenience, I split the **lib/non_local.py** into four python files, and move the 
+old versions (**lib/non_loca.py** and **lib/non_local_simple_version.py**) into 
+**lib/backup/**.
+
+
+## Running Steps
+1. Select the type of non-local block in **lib/network.py**.
+    ```
+    from lib.non_local_concatenation import NONLocalBlock2D
+    from lib.non_local_gaussian import NONLocalBlock2D
+    from lib.non_local_embedded_gaussian import NONLocalBlock2D
+    from lib.non_local_dot_product import NONLocalBlock2D
+2. Run **demo_MNIST.py** with one GPU or multi GPU.
+    ```
+    CUDA_VISIBLE_DEVICES=0,1 python demo_MNIST.py
+
 ## Todo
 - Experiments on Kinetics dataset.
 - 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 (thanks for the person who shares the reason, you can find it in [here](https://github.com/pytorch/pytorch/issues/8637)).

config.py

@@ -1,31 +1,27 @@
-import logging
 import torch
 import torch.utils.data as Data
 import torchvision
 from lib.network import Network
 from torch.autograd import Variable
 from torch import nn
-import config as cfg
-from lib.utils import create_architecture
+import time
 
 
 def calc_acc(x, y):
     x = torch.max(x, dim=-1)[1]
     accuracy = sum(x == y) / x.size(0)
     return accuracy
 
-logging.getLogger().setLevel(logging.INFO)
-
-create_architecture()
 
 train_data = torchvision.datasets.MNIST(root='./mnist', train=True,
                                         transform=torchvision.transforms.ToTensor(),
                                         download=True)
 test_data = torchvision.datasets.MNIST(root='./mnist/',
                                        transform=torchvision.transforms.ToTensor(),
                                        train=False)
-train_loader = Data.DataLoader(dataset=train_data, batch_size=cfg.batch_size, shuffle=True)
-test_loader = Data.DataLoader(dataset=test_data, batch_size=cfg.batch_size, shuffle=False)
+
+train_loader = Data.DataLoader(dataset=train_data, batch_size=128, shuffle=True)
+test_loader = Data.DataLoader(dataset=test_data, batch_size=128, shuffle=False)
 
 train_batch_num = len(train_loader)
 test_batch_num = len(test_loader)
@@ -35,20 +31,19 @@ def calc_acc(x, y):
     net = nn.DataParallel(net)
     net.cuda()
 
-opt = torch.optim.Adam(net.parameters(), lr=cfg.LR, weight_decay=cfg.weight_decay)
+opt = torch.optim.Adam(net.parameters(), lr=0.001)
 loss_func = nn.CrossEntropyLoss()
 
-if cfg.load_model:
-    net.load_state_dict(torch.load(cfg.model_path))
 
-for epoch_index in range(cfg.epoch):
+for epoch_index in range(20):
+    st = time.time()
     for train_batch_index, (img_batch, label_batch) in enumerate(train_loader):
         img_batch = Variable(img_batch)
         label_batch = Variable(label_batch)
 
         if torch.cuda.is_available():
-            img_batch = img_batch.cuda(cfg.cuda_num)
-            label_batch = label_batch.cuda(cfg.cuda_num)
+            img_batch = img_batch.cuda()
+            label_batch = label_batch.cuda()
 
         predict = net(img_batch)
         acc = calc_acc(predict.cpu().data, label_batch.cpu().data)
@@ -58,15 +53,9 @@ def calc_acc(x, y):
         loss.backward()
         opt.step()
 
-        # logging.info('epoch[%d/%d] batch[%d/%d] loss:%.4f acc:%.4f' %
-        #              (epoch_index, cfg.epoch, train_batch_index, train_batch_num, loss.data[0], acc))
-
-    opt.param_groups[0]['lr'] = cfg.LR * (cfg.LR_decay_rate ** (epoch_index // cfg.LR_decay_every_epoch))
-    print('LR', opt.param_groups[0]['lr'])
-# if (train_batch_index + 1) % cfg.test_per_batch == 0:
+    print('(LR:%f) Time of a epoch:%.4fs' % (opt.param_groups[0]['lr'], time.time()-st))
 
     net.eval()
-
     total_loss = 0
     total_acc = 0
 
@@ -75,8 +64,8 @@ def calc_acc(x, y):
         label_batch = Variable(label_batch, volatile=True)
 
         if torch.cuda.is_available():
-            img_batch = img_batch.cuda(cfg.cuda_num)
-            label_batch = label_batch.cuda(cfg.cuda_num)
+            img_batch = img_batch.cuda()
+            label_batch = label_batch.cuda()
 
         predict = net(img_batch)
         acc = calc_acc(predict.cpu().data, label_batch.cpu().data)
@@ -89,7 +78,6 @@ def calc_acc(x, y):
 
     mean_acc = total_acc / test_batch_num
     mean_loss = total_loss / test_batch_num
-    logging.info('[Test] epoch[%d/%d] acc:%.4f loss:%.4f '
-                 % (epoch_index, cfg.epoch, mean_acc, mean_loss.data[0]))
 
-    torch.save(net.state_dict(), cfg.model_path)
+    print('[Test] epoch[%d/%d] acc:%.4f loss:%.4f\n'
+          % (epoch_index, 100, mean_acc, mean_loss.data[0]))

demo_MNIST.py

@@ -27,15 +27,15 @@ def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded
 
         if dimension == 3:
             conv_nd = nn.Conv3d
-            max_pool = nn.MaxPool3d
+            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
             bn = nn.BatchNorm3d
         elif dimension == 2:
             conv_nd = nn.Conv2d
-            max_pool = nn.MaxPool2d
+            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
             bn = nn.BatchNorm2d
         else:
             conv_nd = nn.Conv1d
-            max_pool = nn.MaxPool1d
+            max_pool_layer = nn.MaxPool1d(kernel_size=(2))
             bn = nn.BatchNorm1d
 
         self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
@@ -91,11 +91,13 @@ def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded
         #     self.operation_function = self._gaussian
 
         if sub_sample:
-            self.g = nn.Sequential(self.g, max_pool(kernel_size=2))
+            self.g = nn.Sequential(self.g, max_pool_layer)
             if self.phi is None:
-                self.phi = max_pool(kernel_size=2)
+                self.phi = max_pool_layer
             else:
-                self.phi = nn.Sequential(self.phi, max_pool(kernel_size=2))
+                self.phi = nn.Sequential(self.phi, max_pool_layer)
+
+        print(self.phi)
 
     def forward(self, x):
         '''
@@ -141,20 +143,25 @@ def _embedded_gaussian(self, x):
 
     def _gaussian(self, x):
         batch_size = x.size(0)
+
         g_x = self.g(x).view(batch_size, self.inter_channels, -1)
+
         g_x = g_x.permute(0, 2, 1)
 
         theta_x = x.view(batch_size, self.in_channels, -1)
         theta_x = theta_x.permute(0, 2, 1)
 
         if self.sub_sample:
+            print(self.phi(x).size())
             phi_x = self.phi(x).view(batch_size, self.in_channels, -1)
         else:
             phi_x = x.view(batch_size, self.in_channels, -1)
-
+        print(phi_x.size())
         f = torch.matmul(theta_x, phi_x)
         f_div_C = F.softmax(f, dim=-1)
 
+        print(f_div_C.size(), g_x.size())
+
         y = torch.matmul(f_div_C, g_x)
         y = y.permute(0, 2, 1).contiguous()
         y = y.view(batch_size, self.inter_channels, *x.size()[2:])
@@ -247,23 +254,23 @@ def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', s
 if __name__ == '__main__':
     from torch.autograd import Variable
 
-    mode_list = ['concatenation', 'embedded_gaussian', 'gaussian', 'dot_product', ]
-    # mode_list = ['concatenation']
+    # mode_list = ['concatenation', 'embedded_gaussian', 'gaussian', 'dot_product', ]
+    mode_list = ['gaussian']
 
     for mode in mode_list:
         print(mode)
-        img = Variable(torch.zeros(2, 4, 5))
-        net = NONLocalBlock1D(4, mode=mode, sub_sample=True)
+        img = Variable(torch.zeros(2, 6, 20))
+        net = NONLocalBlock1D(6, mode=mode, sub_sample=True)
         out = net(img)
         print(out.size())
 
-        img = Variable(torch.zeros(2, 4, 10, 10))
-        net = NONLocalBlock2D(4, mode=mode, sub_sample=False, bn_layer=False)
-        out = net(img)
-        print(out.size())
-
-        img = Variable(torch.zeros(2, 4, 5, 4, 5))
-        net = NONLocalBlock3D(4, mode=mode)
-        out = net(img)
-        print(out.size())
+        # img = Variable(torch.zeros(2, 4, 20, 20))
+        # net = NONLocalBlock2D(4, mode=mode, sub_sample=False, bn_layer=False)
+        # out = net(img)
+        # print(out.size())
+        #
+        # img = Variable(torch.zeros(2, 4, 10, 20, 20))
+        # net = NONLocalBlock3D(4, mode=mode)
+        # out = net(img)
+        # print(out.size())
 

lib/__pycache__/non_local_simple_version.cpython-35.pyc

@@ -22,15 +22,15 @@ def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=Tru
 
         if dimension == 3:
             conv_nd = nn.Conv3d
-            max_pool = nn.MaxPool3d
+            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
             bn = nn.BatchNorm3d
         elif dimension == 2:
             conv_nd = nn.Conv2d
-            max_pool = nn.MaxPool2d
+            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
             bn = nn.BatchNorm2d
         else:
             conv_nd = nn.Conv1d
-            max_pool = nn.MaxPool1d
+            max_pool_layer = nn.MaxPool1d(kernel_size=(2))
             bn = nn.BatchNorm1d
 
         self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
@@ -56,8 +56,8 @@ def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=Tru
                            kernel_size=1, stride=1, padding=0)
 
         if sub_sample:
-            self.g = nn.Sequential(self.g, max_pool(kernel_size=2))
-            self.phi = nn.Sequential(self.phi, max_pool(kernel_size=2))
+            self.g = nn.Sequential(self.g, max_pool_layer)
+            self.phi = nn.Sequential(self.phi, max_pool_layer)
 
     def forward(self, x):
         '''
@@ -112,20 +112,22 @@ def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=T
 if __name__ == '__main__':
     from torch.autograd import Variable
     import torch
-    sub_sample = False
 
-    img = Variable(torch.zeros(2, 4, 5))
-    net = NONLocalBlock1D(4, sub_sample=sub_sample, bn_layer=False)
+    sub_sample = True
+    bn_layer = True
+
+    img = Variable(torch.zeros(2, 3, 20))
+    net = NONLocalBlock1D(3, sub_sample=sub_sample, bn_layer=bn_layer)
     out = net(img)
     print(out.size())
 
-    img = Variable(torch.zeros(2, 4, 5, 3))
-    net = NONLocalBlock2D(4, sub_sample=sub_sample)
+    img = Variable(torch.zeros(2, 3, 20, 20))
+    net = NONLocalBlock2D(3, sub_sample=sub_sample, bn_layer=bn_layer)
     out = net(img)
     print(out.size())
 
-    img = Variable(torch.zeros(2, 4, 5, 4, 5))
-    net = NONLocalBlock3D(4, sub_sample=sub_sample)
+    img = Variable(torch.randn(2, 3, 10, 20, 20))
+    net = NONLocalBlock3D(3, sub_sample=sub_sample, bn_layer=bn_layer)
     out = net(img)
     print(out.size())
 

lib/non_local.py lib/backup/non_local.pylib/non_local.py renamed to lib/backup/non_local.py

@@ -1,6 +1,8 @@
 from torch import nn
-from lib.non_local_simple_version import NONLocalBlock2D
-# from lib.non_local import NONLocalBlock2D
+# from lib.non_local_concatenation import NONLocalBlock2D
+# from lib.non_local_gaussian import NONLocalBlock2D
+from lib.non_local_embedded_gaussian import NONLocalBlock2D
+# from lib.non_local_dot_product import NONLocalBlock2D
 
 
 class Network(nn.Module):