modify to support visualize NL_MAP and based on pytorch 1.2.0

Author: AlexHex7

README.md

@@ -3,40 +3,68 @@
 
 ## Statement
 - 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!
 
+- You can **visualize** the Non_local Attention Map by following the **Running Steps** shown below.
+
+- 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!
 
 ## Environment
 - python 3.7.3
 - pytorch 1.2.0
+- opencv 3.4.2
+
+## Visualization
+1. In the **first** Non-local Layer.
+![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_1/37.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_1/44.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_1/46.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_1/110.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_1/161.png)
+
+2. In the **second** Non-local Layer.
+![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_2/1.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_2/8.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_2/10.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_2/18.png)![](/home/phor/Code/project/Non-Local_pytorch/nl_map_vis/nl_map_2/38.png)
+
+
+## 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_train.py** with one GPU or multi GPU to train the Network. Then the weights will be save in **weights/**.
+    ```
+    CUDA_VISIBLE_DEVICES=0,1 python demo_MNIST.py
+
+3. Run **nl_map_save.py** to save NL_MAP of one test sample in **nl_map_vis**.
+    ```
+    CUDA_VISIBLE_DEVICES=0,1 python nl_map_save.py
+    
+4. Come into **nl_map_vis/** and run **nl_map_vis.py** to visualize the NL_MAP. (tips: if the Non-local type you select is **non_local_concatenation** or **non_local_dot_product** (without Softmax operation), you may need to normalize NL_MAP in the visualize code)
+    ```
+    python nl_map_save.py
 
 
 ## 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 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/**.
+
 6. Test the code with pytorch 1.1.0 and it works.
-7. Move the code supporting pytorch 0.4.1 and 1.1.0 to **Non-Local_pytorch_0.4.1_to_1.1.0/**.
 
-## 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
+7. Move the code supporting pytorch 0.4.1 and 1.1.0 to **Non-Local_pytorch_0.4.1_to_1.1.0/** (In fact, I think it can also support pytorch 1.2.0).
+
+8. In order to visualize NL_MAP, some code have been slightly modified. The code **nl_map_save.py** is added to save NL_MAP (two Non-local Layer) of one test sample. The code **Non-local_pytorch/nl_map_vis.py** is added to visualize NL_MAP. Besieds, the code is support pytorch 1.2.0.
+
 
 ## Todo
 - Experiments on Charades dataset.

demo_MNIST_train.py

@@ -0,0 +1,81 @@
+import torch
+import torch.utils.data as Data
+import torchvision
+from lib.network import Network
+from torch import nn
+import time
+
+
+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(10):
+    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)
+        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, 10, mean_acc * 100, mean_loss.item()))
+
+# weight_path = 'weights/net.pth'
+# print('Save Net weights to', weight_path)
+# net.cpu()
+# torch.save(net.state_dict(), weight_path)

lib/network.py

@@ -0,0 +1,79 @@
+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.conv_1 = 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),
+        )
+
+        self.nl_1 = NONLocalBlock2D(in_channels=32)
+        self.conv_2 = nn.Sequential(
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+        )
+
+        self.nl_2 = NONLocalBlock2D(in_channels=64)
+        self.conv_3 = nn.Sequential(
+            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)
+
+        feature_1 = self.conv_1(x)
+        nl_feature_1 = self.nl_1(feature_1)
+
+        feature_2 = self.conv_2(nl_feature_1)
+        nl_feature_2 = self.nl_2(feature_2)
+
+        output = self.conv_3(nl_feature_2).view(batch_size, -1)
+        output = self.fc(output)
+
+        return output
+
+    def forward_with_nl_map(self, x):
+        batch_size = x.size(0)
+
+        feature_1 = self.conv_1(x)
+        nl_feature_1, nl_map_1 = self.nl_1(feature_1, return_nl_map=True)
+
+        feature_2 = self.conv_2(nl_feature_1)
+        nl_feature_2, nl_map_2 = self.nl_2(feature_2, return_nl_map=True)
+
+        output = self.conv_3(nl_feature_2).view(batch_size, -1)
+        output = self.fc(output)
+
+        return output, [nl_map_1, nl_map_2]
+
+
+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,150 @@
+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, return_nl_map=False):
+        '''
+        :param x: (b, c, t, h, w)
+        :param return_nl_map: if True return z, nl_map, else only return z.
+        :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
+
+        if return_nl_map:
+            return z, f_div_C
+        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())