code supporting pytorch 0.4.1

Author: AlexHex7

README.md

@@ -10,17 +10,19 @@ in **lib/network.py**.
 
 ## Environment
 - python 3.6
-- pytorch 0.3.0
+- pytorch 0.4.1
 
 ## 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 
+3. Fix the error 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/**.
+5. modify the code to support pytorch 0.4.1, and move the code supporting pytorch 0.3.1 \
+to **Non-Local_pytorch_0.3.1/**.
 
 
 ## Running Steps

demo_MNIST.py

@@ -0,0 +1,83 @@
+import torch
+import torch.utils.data as Data
+import torchvision
+from lib.network import Network
+from torch import nn
+import time
+
+
+def calc_acc(x, y):
+    x = torch.max(x, dim=-1)[1]
+    accuracy = sum(x == y) / x.size(0)
+    return accuracy
+
+
+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=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)
+
+net = Network()
+if torch.cuda.is_available():
+    net = nn.DataParallel(net)
+    net.cuda()
+
+opt = torch.optim.Adam(net.parameters(), lr=0.001)
+loss_func = nn.CrossEntropyLoss()
+
+for epoch_index in range(20):
+    st = time.time()
+
+    torch.set_grad_enabled(True)
+    net.train()
+    for train_batch_index, (img_batch, label_batch) in enumerate(train_loader):
+        if torch.cuda.is_available():
+            img_batch = img_batch.cuda()
+            label_batch = label_batch.cuda()
+
+        predict = net(img_batch)
+        acc = calc_acc(predict.cpu().data, label_batch.cpu().data)
+        loss = loss_func(predict, label_batch)
+
+        net.zero_grad()
+        loss.backward()
+        opt.step()
+
+    print('(LR:%f) Time of a epoch:%.4fs' % (opt.param_groups[0]['lr'], time.time()-st))
+
+    torch.set_grad_enabled(False)
+    net.eval()
+    total_loss = []
+    total_acc = 0
+    total_sample = 0
+
+    for test_batch_index, (img_batch, label_batch) in enumerate(test_loader):
+        if torch.cuda.is_available():
+            img_batch = img_batch.cuda()
+            label_batch = label_batch.cuda()
+
+        predict = net(img_batch)
+        loss = loss_func(predict, label_batch)
+
+        predict = predict.argmax(dim=1)
+        acc = (predict == label_batch).sum()
+
+        total_loss.append(loss)
+        total_acc += acc
+        total_sample += img_batch.size(0)
+
+    net.train()
+
+    mean_acc = total_acc.item() * 1.0 / total_sample
+    mean_loss = sum(total_loss) / total_loss.__len__()
+
+    print('[Test] epoch[%d/%d] acc:%.4f%% loss:%.4f\n'
+          % (epoch_index, 100, mean_acc * 100, mean_loss.item()))

lib/network.py

@@ -0,0 +1,52 @@
+from torch import nn
+# 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):
+    def __init__(self):
+        super(Network, self).__init__()
+
+        self.convs = nn.Sequential(
+            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+
+            NONLocalBlock2D(in_channels=32),
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+
+            NONLocalBlock2D(in_channels=64),
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+        )
+
+        self.fc = nn.Sequential(
+            nn.Linear(in_features=128*3*3, out_features=256),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+
+            nn.Linear(in_features=256, out_features=10)
+        )
+
+    def forward(self, x):
+        batch_size = x.size(0)
+        output = self.convs(x).view(batch_size, -1)
+        output = self.fc(output)
+        return output
+
+if __name__ == '__main__':
+    import torch
+
+    img = torch.randn(3, 1, 28, 28)
+    net = Network()
+    out = net(img)
+    print(out.size())
+

lib/non_local_concatenation.py

@@ -0,0 +1,147 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class _NonLocalBlockND(nn.Module):
+    def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):
+        super(_NonLocalBlockND, self).__init__()
+
+        assert dimension in [1, 2, 3]
+
+        self.dimension = dimension
+        self.sub_sample = sub_sample
+
+        self.in_channels = in_channels
+        self.inter_channels = inter_channels
+
+        if self.inter_channels is None:
+            self.inter_channels = in_channels // 2
+            if self.inter_channels == 0:
+                self.inter_channels = 1
+
+        if dimension == 3:
+            conv_nd = nn.Conv3d
+            max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))
+            bn = nn.BatchNorm3d
+        elif dimension == 2:
+            conv_nd = nn.Conv2d
+            max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))
+            bn = nn.BatchNorm2d
+        else:
+            conv_nd = nn.Conv1d
+            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,
+                         kernel_size=1, stride=1, padding=0)
+
+        if bn_layer:
+            self.W = nn.Sequential(
+                conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
+                        kernel_size=1, stride=1, padding=0),
+                bn(self.in_channels)
+            )
+            nn.init.constant_(self.W[1].weight, 0)
+            nn.init.constant_(self.W[1].bias, 0)
+        else:
+            self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
+                             kernel_size=1, stride=1, padding=0)
+            nn.init.constant_(self.W.weight, 0)
+            nn.init.constant_(self.W.bias, 0)
+
+        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)
+
+        self.concat_project = nn.Sequential(
+            nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
+            nn.ReLU()
+        )
+
+        if sub_sample:
+            self.g = nn.Sequential(self.g, max_pool_layer)
+            self.phi = nn.Sequential(self.phi, max_pool_layer)
+
+    def forward(self, x):
+        '''
+        :param x: (b, c, t, h, w)
+        :return:
+        '''
+
+        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)
+
+        # (b, c, N, 1)
+        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1, 1)
+        # (b, c, 1, N)
+        phi_x = self.phi(x).view(batch_size, self.inter_channels, 1, -1)
+
+        h = theta_x.size(2)
+        w = phi_x.size(3)
+        theta_x = theta_x.repeat(1, 1, 1, w)
+        phi_x = phi_x.repeat(1, 1, h, 1)
+
+        concat_feature = torch.cat([theta_x, phi_x], dim=1)
+        f = self.concat_project(concat_feature)
+        b, _, h, w = f.size()
+        f = f.view(b, h, w)
+
+        N = f.size(-1)
+        f_div_C = f / N
+
+        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:])
+        W_y = self.W(y)
+        z = W_y + x
+
+        return z
+
+
+class NONLocalBlock1D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
+        super(NONLocalBlock1D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=1, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+class NONLocalBlock2D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
+        super(NONLocalBlock2D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=2, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+class NONLocalBlock3D(_NonLocalBlockND):
+    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
+        super(NONLocalBlock3D, self).__init__(in_channels,
+                                              inter_channels=inter_channels,
+                                              dimension=3, sub_sample=sub_sample,
+                                              bn_layer=bn_layer)
+
+
+if __name__ == '__main__':
+    import torch
+
+    for (sub_sample, bn_layer) in [(True, True), (False, False), (True, False), (False, True)]:
+        img = torch.zeros(2, 3, 20)
+        net = NONLocalBlock1D(3, sub_sample=sub_sample, bn_layer=bn_layer)
+        out = net(img)
+        print(out.size())
+
+        img = torch.zeros(2, 3, 20, 20)
+        net = NONLocalBlock2D(3, sub_sample=sub_sample, bn_layer=bn_layer)
+        out = net(img)
+        print(out.size())
+
+        img = torch.randn(2, 3, 8, 20, 20)
+        net = NONLocalBlock3D(3, sub_sample=sub_sample, bn_layer=bn_layer)
+        out = net(img)
+        print(out.size())