Dynamic conv test

configs/train-div2k/train_edsr-baseline-liif.yaml

@@ -2,7 +2,7 @@ train_dataset:
   dataset:
     name: image-folder
     args:
-      root_path: ./load/div2k/DIV2K_train_HR
+      root_path: /data/home/wanghanying/work/dataset/DIV2K/DIV2K_train_HR
       repeat: 20
       cache: in_memory
   wrapper:
@@ -18,7 +18,7 @@ val_dataset:
   dataset:
     name: image-folder
     args:
-      root_path: ./load/div2k/DIV2K_valid_HR
+      root_path: /data/home/wanghanying/work/dataset/DIV2K/DIV2K_valid_HR
       first_k: 10
       repeat: 160
       cache: in_memory
@@ -58,3 +58,4 @@ multi_step_lr:
 
 epoch_val: 1
 epoch_save: 100
+# resume: /data/home/wanghanying/work/liif/save/train_edsr_dytest/epoch-last.pth

models/dynamic_conv.py

@@ -0,0 +1,342 @@
+#https://github.com/kaijieshi7/Dynamic-convolution-Pytorch/blob/master/dynamic_conv.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+
+class attention1d(nn.Module):
+    def __init__(self, in_planes, ratios, K, temperature, init_weight=True):
+        super(attention1d, self).__init__()
+        assert temperature%3==1
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        if in_planes!=3:
+            hidden_planes = int(in_planes*ratios)+1
+        else:
+            hidden_planes = K
+        self.fc1 = nn.Conv1d(in_planes, hidden_planes, 1, bias=False)
+        # self.bn = nn.BatchNorm2d(hidden_planes)
+        self.fc2 = nn.Conv1d(hidden_planes, K, 1, bias=True)
+        self.temperature = temperature
+        if init_weight:
+            self._initialize_weights()
+
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            if isinstance(m ,nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def updata_temperature(self):
+        if self.temperature!=1:
+            self.temperature -=3
+            print('Change temperature to:', str(self.temperature))
+
+
+    def forward(self, x):
+        x = self.avgpool(x)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x).view(x.size(0), -1)
+        return F.softmax(x/self.temperature, 1)
+
+
+class Dynamic_conv1d(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, ratio=0.25, stride=1, padding=0, dilation=1, groups=1, bias=True, K=4,temperature=34, init_weight=True):
+        super(Dynamic_conv1d, self).__init__()
+        assert in_planes%groups==0
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.bias = bias
+        self.K = K
+        self.attention = attention1d(in_planes, ratio, K, temperature)
+
+        self.weight = nn.Parameter(torch.randn(K, out_planes, in_planes//groups, kernel_size), requires_grad=True)
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(K, out_planes))
+        else:
+            self.bias = None
+        if init_weight:
+            self._initialize_weights()
+
+        #TODO 初始化
+    def _initialize_weights(self):
+        for i in range(self.K):
+            nn.init.kaiming_uniform_(self.weight[i])
+
+
+    def update_temperature(self):
+        self.attention.updata_temperature()
+
+    def forward(self, x):#将batch视作维度变量，进行组卷积，因为组卷积的权重是不同的，动态卷积的权重也是不同的
+        softmax_attention = self.attention(x)
+        batch_size, in_planes, height = x.size()
+        x = x.view(1, -1, height, )# 变化成一个维度进行组卷积
+        weight = self.weight.view(self.K, -1)
+
+        # 动态卷积的权重的生成， 生成的是batch_size个卷积参数（每个参数不同）
+        aggregate_weight = torch.mm(softmax_attention, weight).view(batch_size*self.out_planes, self.in_planes//self.groups, self.kernel_size,)
+        if self.bias is not None:
+            aggregate_bias = torch.mm(softmax_attention, self.bias).view(-1)
+            output = F.conv1d(x, weight=aggregate_weight, bias=aggregate_bias, stride=self.stride, padding=self.padding,
+                              dilation=self.dilation, groups=self.groups*batch_size)
+        else:
+            output = F.conv1d(x, weight=aggregate_weight, bias=None, stride=self.stride, padding=self.padding,
+                              dilation=self.dilation, groups=self.groups * batch_size)
+
+        output = output.view(batch_size, self.out_planes, output.size(-1))
+        return output
+
+
+
+class attention2d(nn.Module):
+    def __init__(self, in_planes, ratios, K, temperature, init_weight=True):
+        super(attention2d, self).__init__()
+        assert temperature%3==1
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        if in_planes!=3:
+            hidden_planes = int(in_planes*ratios)+1
+        else:
+            hidden_planes = K
+        self.fc1 = nn.Conv2d(in_planes, hidden_planes, 1, bias=False)
+        # self.bn = nn.BatchNorm2d(hidden_planes)
+        self.fc2 = nn.Conv2d(hidden_planes, K, 1, bias=True)
+        self.temperature = temperature
+        if init_weight:
+            self._initialize_weights()
+
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            if isinstance(m ,nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def updata_temperature(self):
+        if self.temperature!=1:
+            self.temperature -=3
+            print('Change temperature to:', str(self.temperature))
+
+
+    def forward(self, x):
+        x = self.avgpool(x)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x).view(x.size(0), -1)
+        return F.softmax(x/self.temperature, 1)
+
+
+class Dynamic_conv2d(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, ratio=0.25, stride=1, padding=0, dilation=1, groups=1, bias=True, K=4,temperature=34, init_weight=True):
+        super(Dynamic_conv2d, self).__init__()
+        assert in_planes%groups==0
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.bias = bias
+        self.K = K
+        self.attention = attention2d(in_planes, ratio, K, temperature)
+
+        self.weight = nn.Parameter(torch.randn(K, out_planes, in_planes//groups, kernel_size, kernel_size), requires_grad=True)
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(K, out_planes))
+        else:
+            self.bias = None
+        if init_weight:
+            self._initialize_weights()
+
+        #TODO 初始化
+    def _initialize_weights(self):
+        for i in range(self.K):
+            nn.init.kaiming_uniform_(self.weight[i])
+
+
+    def update_temperature(self):
+        self.attention.updata_temperature()
+
+    def forward(self, x):#将batch视作维度变量，进行组卷积，因为组卷积的权重是不同的，动态卷积的权重也是不同的
+        softmax_attention = self.attention(x)
+        batch_size, in_planes, height, width = x.size()
+        x = x.view(1, -1, height, width)# 变化成一个维度进行组卷积
+        weight = self.weight.view(self.K, -1)
+
+        # 动态卷积的权重的生成， 生成的是batch_size个卷积参数（每个参数不同）
+        aggregate_weight = torch.mm(softmax_attention, weight).view(batch_size*self.out_planes, self.in_planes//self.groups, self.kernel_size, self.kernel_size)
+        if self.bias is not None:
+            aggregate_bias = torch.mm(softmax_attention, self.bias).view(-1)
+            output = F.conv2d(x, weight=aggregate_weight, bias=aggregate_bias, stride=self.stride, padding=self.padding,
+                              dilation=self.dilation, groups=self.groups*batch_size)
+        else:
+            output = F.conv2d(x, weight=aggregate_weight, bias=None, stride=self.stride, padding=self.padding,
+                              dilation=self.dilation, groups=self.groups * batch_size)
+
+        output = output.view(batch_size, self.out_planes, output.size(-2), output.size(-1))
+        return output
+
+
+class attention3d(nn.Module):
+    def __init__(self, in_planes, ratios, K, temperature):
+        super(attention3d, self).__init__()
+        assert temperature%3==1
+        self.avgpool = nn.AdaptiveAvgPool3d(1)
+        if in_planes != 3:
+            hidden_planes = int(in_planes * ratios)+1
+        else:
+            hidden_planes = K
+        self.fc1 = nn.Conv3d(in_planes, hidden_planes, 1, bias=False)
+        self.fc2 = nn.Conv3d(hidden_planes, K, 1, bias=False)
+        self.temperature = temperature
+
+    def updata_temperature(self):
+        if self.temperature!=1:
+            self.temperature -=3
+            print('Change temperature to:', str(self.temperature))
+
+    def forward(self, x):
+        x = self.avgpool(x)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x).view(x.size(0), -1)
+        return F.softmax(x / self.temperature, 1)
+
+class Dynamic_conv3d(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, ratio=0.25, stride=1, padding=0, dilation=1, groups=1, bias=True, K=4, temperature=34):
+        super(Dynamic_conv3d, self).__init__()
+        assert in_planes%groups==0
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.bias = bias
+        self.K = K
+        self.attention = attention3d(in_planes, ratio, K, temperature)
+
+        self.weight = nn.Parameter(torch.randn(K, out_planes, in_planes//groups, kernel_size, kernel_size, kernel_size), requires_grad=True)
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(K, out_planes))
+        else:
+            self.bias = None
+
+
+        #TODO 初始化
+        # nn.init.kaiming_uniform_(self.weight, )
+
+    def update_temperature(self):
+        self.attention.updata_temperature()
+
+    def forward(self, x):#将batch视作维度变量，进行组卷积，因为组卷积的权重是不同的，动态卷积的权重也是不同的
+        softmax_attention = self.attention(x)
+        batch_size, in_planes, depth, height, width = x.size()
+        x = x.view(1, -1, depth, height, width)# 变化成一个维度进行组卷积
+        weight = self.weight.view(self.K, -1)
+
+        # 动态卷积的权重的生成， 生成的是batch_size个卷积参数（每个参数不同）
+        aggregate_weight = torch.mm(softmax_attention, weight).view(batch_size*self.out_planes, self.in_planes//self.groups, self.kernel_size, self.kernel_size, self.kernel_size)
+        if self.bias is not None:
+            aggregate_bias = torch.mm(softmax_attention, self.bias).view(-1)
+            output = F.conv3d(x, weight=aggregate_weight, bias=aggregate_bias, stride=self.stride, padding=self.padding,
+                              dilation=self.dilation, groups=self.groups*batch_size)
+        else:
+            output = F.conv3d(x, weight=aggregate_weight, bias=None, stride=self.stride, padding=self.padding,
+                              dilation=self.dilation, groups=self.groups * batch_size)
+
+        output = output.view(batch_size, self.out_planes, output.size(-3), output.size(-2), output.size(-1))
+        return output
+
+
+def conv3x3(in_planes, out_planes, kernel_size=3, bias=False, stride=1, groups=1, dilation=1):
+    return Dynamic_conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=bias, dilation=dilation)
+
+# def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+#     return Dynamic_conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+# def conv1x1(in_planes, out_planes, stride=1):
+#     """1x1 convolution"""
+#     return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+def conv1x1(in_planes, out_planes, stride=1):
+    return Dynamic_conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False,)
+
+
+if __name__ == '__main__':
+    x = torch.randn(24, 3,  20)
+    model = Dynamic_conv1d(in_planes=3, out_planes=16, kernel_size=3, ratio=0.25, padding=1,)
+    x = x.to('cuda:0')
+    model.to('cuda')
+    # model.attention.cuda()
+    # nn.Conv3d()
+    print(model(x).shape)
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    model.update_temperature()
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+    print(model(x).shape)
+