LeNet模型

原理框架图


代码实现

首先定义我们的LeNet模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import torch
from torch import nn
from d2l import torch as d2l


class Reshape(torch.nn.Module):
def forward(self, x):
return x.view(-1, 1, 28, 28)


# 构建网络
net = torch.nn.Sequential(
Reshape(), # 对应 阶段1
nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.Sigmoid(), # 对应 阶段2
nn.AvgPool2d(kernel_size=2, stride=2), # 对应 阶段3
nn.Conv2d(6, 16, kernel_size=5), nn.Sigmoid(), # 对应 阶段4
nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(), # 对应 阶段5
nn.Linear(16*5*5, 120), nn.Sigmoid(), # 对应 阶段6
nn.Linear(120, 84), nn.Sigmoid(), # 对应 阶段7
nn.Linear(84, 10) # 对应 阶段8
)
# torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None)

检验模型

1
2
3
4
5
6
X = torch.rand(size=(1, 1, 28, 28), dtype=torch.float32)

for layer in net:
X = layer(X)
print(layer.__class__.__name__, 'output shape:\t', X.shape)

Reshape output shape:     torch.Size([1, 1, 28, 28])
Conv2d output shape:     torch.Size([1, 6, 28, 28])
Sigmoid output shape:     torch.Size([1, 6, 28, 28])
AvgPool2d output shape:     torch.Size([1, 6, 14, 14])
Conv2d output shape:     torch.Size([1, 16, 10, 10])
Sigmoid output shape:     torch.Size([1, 16, 10, 10])
AvgPool2d output shape:     torch.Size([1, 16, 5, 5])
Flatten output shape:     torch.Size([1, 400])
Linear output shape:     torch.Size([1, 120])
Sigmoid output shape:     torch.Size([1, 120])
Linear output shape:     torch.Size([1, 84])
Sigmoid output shape:     torch.Size([1, 84])
Linear output shape:     torch.Size([1, 10])