PyTorch 自定义层（区别于自定义模型）

博主在学习三值神经网络时，使用了LeNet-5模型，编程代码，需要对LeNet-5模型中的卷积层与全连接层进行自定义，搜索他人方法后，博主产生了一个疑问，绝大多数提供的自定义层方法都是继承 nn.Module 模型，而这方法据说是官方提供（官网网址：PyTorch），自定义线性层代码如下：

class Linear(nn.Module):
    def __init__(self, input_features, output_features, bias=True):
        super(Linear, self).__init__()
        self.input_features = input_features
        self.output_features = output_features
        self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(output_features))
        else:
            self.register_parameter('bias', None)
        self.weight.data.uniform_(-0.1, 0.1)
        if bias is not None:
            self.bias.data.uniform_(-0.1, 0.1)
    def forward(self, input):
        return LinearFunction.apply(input, self.weight, self.bias)

博主因涉猎PyTorch时长过短，对此了解不多，初始仅是觉得奇怪，为何自定义层需要继承nn.Module模型，乍一看像是自定义模型而非自定义层

因博主过于小白，当时就想，这么误人子弟的自定义方法不应该笼统地用nn.Module一齐实现，PyTorch是一个成熟的软件，应该知道还有博主这种小小白没办法立刻上手以nn.Module为依赖的自定义层方法（说实话，博主看了几天才看懂上面的代码是啥意思……哭泣……）

所以博主又搜索了GitHub中与已发表论文相关的PyTorch深度学习代码，终于发现了PyTorch自定义层的简单方法，可以细节到区分具体层的种类，而不会产生混淆

以下为博主使用的LeNet-5模型，为直观体现自定义层的方法，博主仅留下了自定义的卷积层与全连接层，即代码中 TernaryConv2d(1, 32, kernel_size=5) 与 TernaryLinear(512, 10) 部分

# LeNet-5模型
class LeNet_5(nn.Module):
    def __init__(self):
        super(LeNet_5, self).__init__()
        self.conv1 = TernaryConv2d(1, 32, kernel_size=5)  # 卷积
        self.fc1 = TernaryLinear(512, 10)  # 全连接

    def forward(self, x):
        x = self.conv1(x)
        x = self.fc1(x)
        return x

以下为博主自定义的卷积层（TernaryConv2d）与全连接层（TernaryLinear）

class TernaryConv2d(nn.Conv2d):
    # def __init__(self, *args, **kwargs):
    # 该方法接受任意个数的参数，其中不指定key的参数会以list形式保存到args变量中，指定key的参数会以dict的形式保存到kwargs变量中
    def __init__(self, *args, **kwargs):
        super(TernaryConv2d, self).__init__(*args, **kwargs)

    def forward(self, input):
        out = F.conv2d(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
        return out

class TernaryLinear(nn.Linear):
    def __init__(self, *args, **kwargs):
        super(TernaryLinear, self).__init__(*args, **kwargs)

    def forward(self, input):
        out = F.linear(input, self.weight, self.bias)
        return out

以卷积层为例，类名可自行定义，但必须继承 nn.Conv2d ，在该层初始化使用时，若无参数设置，将会默认卷积层的参数，默认参数可查阅官方文档（在 Python API 中）获知，这里博主将其粘贴了过来：

CLASS torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')

而后，在 def init() 函数中可初始化参数，在 def forward() 函数中可对卷积层自行设计

若想自定义其他类型的层，只需修改继承的 nn.Conv2d 即可，操作非常简单且不会产生歧义

博主将模型与层的代码整合如下，便于复制使用：

class TernaryConv2d(nn.Conv2d):
    # def __init__(self, *args, **kwargs):
    # 该方法接受任意个数的参数，其中不指定key的参数会以list形式保存到args变量中，指定key的参数会以dict的形式保存到kwargs变量中
    def __init__(self, *args, **kwargs):
        super(TernaryConv2d, self).__init__(*args, **kwargs)

    def forward(self, input):
        out = F.conv2d(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
        return out


class TernaryLinear(nn.Linear):
    def __init__(self, *args, **kwargs):
        super(TernaryLinear, self).__init__(*args, **kwargs)

    def forward(self, input):
        out = F.linear(input, self.weight, self.bias)
        return out


# LeNet-5模型
class LeNet_5(nn.Module):
    def __init__(self):
        super(LeNet_5, self).__init__()
        self.conv1 = TernaryConv2d(1, 32, kernel_size=5)  # 卷积
        self.fc1 = TernaryLinear(512, 10)  # 全连接

    def forward(self, x):
        x = self.conv1(x)
        x = self.fc1(x)
        return x

因官方模型类API过多，博主无法一一查看，博主猜测nn.Module模型中应该与自定义层所使用的nn.Conv2d有着某种继承关系，所以可以如此使用