pytorch学习笔记 —— torch.nn.Module

torch.nn.Module（以下简称 Module）是所有神经网络模块的基类，在 pytorch 中，自定义层、自定义块、自定义模型都可以通过继承 Module 类来实现；

Module 中有许多方法，在自定义类时必须重写其中的两个方法：__init__ 和 forward；

• __init__ 中存放模型的固有属性，如：全连接层、卷积层等等；
• forward 中写各层之间的连接计算关系，即前向传播（实现模型的功能）；
• 注：有可学习参数的层放在 __init__ 中，无可学习参数的层放在 __init__ 或 forward 中均可；

继承 Module 类来实现一个简单的手写数字识别的 CNN 网络如下：

from torch import nn


# 输入图片大小(1 * 28 * 28)
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 5, 1, 2),  # if stride = 1, padding = (kernel_size - 1) / 2
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output