pytorch 自定义网络层

有时需要自定义卷积核的权重，但torch.nn.Conv2d和torch.nn.Conv3d中的卷积核参数不允许自定义，因此需要使用torch.nn.functional中的Conv2d或nn.Conv3d

以自定义SRM层为例

空间富模型(SRM, Spatial Rich Model)过滤器层使用3个不同参数的卷积核分别提取3种不同的高频残差信号，并将其参数设置为不可训练。SRM过滤器所使用的3个卷积核参数分别为:

实际上是自定义了一个参数不可训练的卷积层。自定义层需要继承 nn.Module 类，并重写 __init__ 和 forward 两个方法。

Conv2d

import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np

class SRM2D(nn.Module):
    def __init__(self):
        super().__init__()

        q = [4.0, 12.0, 4.0]
        filter1 = [[0,  0,  0,  0, 0],
                   [0, -1,  2, -1, 0],
                   [0,  2, -4,  2, 0],
                   [0, -1,  2, -1, 0],
                   [0,  0,  0,  0, 0]]

        filter2 = [[-1, 2, -2,  2,-1],
                   [ 2,-6,  8, -6, 2],
                   [-2, 8, -12, 8,-2],
                   [ 2,-6,  8, -6, 2],
                   [-1, 2, -2,  2,-1]]

        filter3 = [[0,  0,  0,  0, 0],
                   [0,  0, -1,  0, 0],
                   [0, -1, +4, -1, 0],
                   [0,  0, -1,  0, 0],
                   [0,  0,  0,  0, 0]]

        filter1 = np.asarray(filter1, dtype=float) / q[0]
        filter2 = np.asarray(filter2, dtype=float) / q[1]
        filter3 = np.asarray(filter3, dtype=float) / q[2]
        
        # 自定义卷积核权重
        self.filter = torch.tensor([[filter1, filter1, filter1], 
                                    [filter2, filter2, filter2], 
                                    [filter3, filter3, filter3]],
                                    dtype=torch.float32)

    def forward(self, input):
        
        def truncate(x):
            neg = ((x + 2) + abs(x + 2)) / 2 - 2
            return -(-neg + 2 + abs(- neg + 2)) / 2 + 2

        result = F.conv2d(input,
                          weight=nn.Parameter(self.filter, requires_grad=False), # 设置为参数不可训练
                          stride=(1, 1, 1),
                          # 因为卷积核大小为5×5，步长为1，若想保持输出和输入大小相等，需设置padding为2
                          padding=(2, 2, 2)) 

        result = truncate(result)
        return result

---

Conv3d

import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np

class SRM3D(nn.Module):
    def __init__(self):
        super().__init__()

        q = [4.0, 12.0, 4.0]
        filter1 = [[0,  0,  0,  0, 0],
                   [0, -1,  2, -1, 0],
                   [0,  2, -4,  2, 0],
                   [0, -1,  2, -1, 0],
                   [0,  0,  0,  0, 0]]

        filter2 = [[-1, 2, -2,  2,-1],
                   [ 2,-6,  8, -6, 2],
                   [-2, 8, -12, 8,-2],
                   [ 2,-6,  8, -6, 2],
                   [-1, 2, -2,  2,-1]]

        filter3 = [[0,  0,  0,  0, 0],
                   [0,  0, -1,  0, 0],
                   [0, -1, +4, -1, 0],
                   [0,  0, -1,  0, 0],
                   [0,  0,  0,  0, 0]]

        filter1 = np.asarray(filter1, dtype=float) / q[0]
        filter2 = np.asarray(filter2, dtype=float) / q[1]
        filter3 = np.asarray(filter3, dtype=float) / q[2]
        
        # 自定义卷积核权重
        filter = torch.tensor([[filter1, filter1, filter1], 
                               [filter2, filter2, filter2], 
                               [filter3, filter3, filter3]],
                               dtype=torch.float32)
        
        # 因为是3D卷积，所以需要扩充维度
        self.filter = torch.unsqueeze(filter, 2)

    def forward(self, input):
        
        def truncate(x):
            neg = ((x + 2) + abs(x + 2)) / 2 - 2
            return -(-neg + 2 + abs(- neg + 2)) / 2 + 2

        result = F.conv3d(input,
                          weight=nn.Parameter(self.filter, requires_grad=False), # 设置为参数不可训练
                          stride=(1, 1, 1),
                          # 因为卷积核大小为5×5，步长为1，若想保持输出和输入大小相等，需设置padding为(0, 2, 2)
                          padding=(0, 2, 2))

        result = truncate(result)
        return result