【Pytorch中的nn.Module模块】

1. nn.Module基本情况介绍

nn.Module模块已经给我们实现了

• Linear
• ReLU
• Sigmoid
• Conv2d
• ConvTransposed2d
• Dropout
• etc.

我们只需要实现init和forward就可以定义一个神经网络

2. Container(nn.Sequential)

.net(x)
可以重复使用相同的网络结构。

self.net = nn.Sequential(
	nn.Conv2d(1,32,5,1,1)
	nn.MaxPool2d(2,2)
	nn.ReLU(True)
	nn.BatchNorm2d(32)
)

3. parameters

import torch
import torch.nn as nn
import torch.nn.functional as F

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 卷积层 1: 输入通道为1（灰度图），输出通道为6，卷积核大小为5x5
        self.conv1 = nn.Conv2d(1, 6, 5)
		#....
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 通过第一个卷积层后，应用ReLU激活函数，然后进行最大池化
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
		#....
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # 所有维度，除了批处理维度
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

# 实例化网络
net = LeNet5()
print(net)#得到网络结构
print(list(net.parameters())[0])#得到第0层的参数结构
print(list(net.named_parameters())[0])
print(dict(net.named_parameters()).items())
# 很方便的使用优化器来管理参数
optimizer = optim.SGD(net.parameters(),lr=1e-3)

4.modules

• modules
  • 所有的节点
• children
  • 直接相连的节点

nn.Sequential 允许您创建一个模块，其中每个子模块按照它们被添加的顺序依次执行。这可以使代码更加简洁和易于理解，特别是对于那些没有分支的直线部分。
以上代码可以修改为：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, 6, 5),  # 卷积层 1
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # 池化层
            nn.Conv2d(6, 16, 5),  # 卷积层 2
            nn.ReLU(),
            nn.MaxPool2d(2, 2)   # 池化层
        )

        self.fc_layers = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),  # 全连接层 1
            nn.ReLU(),
            nn.Linear(120, 84),  # 全连接层 2
            nn.ReLU(),
            nn.Linear(84, 10)    # 全连接层 3
        )

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(-1, self.num_flat_features(x))
        x = self.fc_layers(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # 所有维度，除了批处理维度
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

5.to(device)

训练放在GPU或者CPU上

device = torch.device('cuda')
net = Net()
net.to(device)

6.save and load

训练参数的加载和保存

device = torch.device('cuda')
net = Net()
net.to(device)
net.load_state_dict(torch.load('ckpt.mdl'))
# train...

torch.save(net.state_dict(),'ckpt.mdl')

7.train/test

训练和测试的模式切换

device = torch.device('cuda')
net = Net()
net.to(device)
net.load_state_dict(torch.load('ckpt.mdl'))
# train...
net.train()
...
# test
net.eval()

8.implement own layer

可以定义自己的层

class Flatten(nn.Module):
	def __init__(self):
		super(Flatten,self).__init__()
	
	def forward(self, input):
		# 将各层之间展平
		return input.view(input.size(0),-1)

class MyLinear(nn.Module):
	def __init__(self):
		super(MyLinear,self).__init__()
		#require_grad=True
		self.w = nn.Parameter(torch.randn(outp, inp))
		self.b = nn.Parameter(torch.randn(outp))
	
	def forward(self, input):
		# 将各层之间展平
		x = x @ self.w.t() + self.b
		return x