LeNet5实现-pytorch

完整代码:PyNet/pytorch/lenet5_test.py

加载数据

pytorch提供模块torchvision,用于数据的加载、预处理和批量化

  • torchvision.datasets内置类MNIST用于mnist数据集下载和加载
  • torchvision.transforms对数据进行预处理
  • torchvision.DataLoader对数据进行批量化
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def load_mnist_data(batch_size=128, shuffle=False):
data_dir = '/home/zj/data/'

transform = transforms.Compose([
transforms.Grayscale(),
transforms.Resize(size=(32, 32)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5,), std=(0.5,))
])

train_data_set = datasets.MNIST(root=data_dir, train=True, download=True, transform=transform)
test_data_set = datasets.MNIST(root=data_dir, train=False, download=True, transform=transform)

train_loader = DataLoader(train_data_set, batch_size=batch_size, shuffle=shuffle)
test_loader = DataLoader(test_data_set, batch_size=batch_size, shuffle=shuffle)

return train_loader, test_loader

网络定义

LeNet-5模型定义参考卷积神经网络推导-单张图片矩阵计算

torch.nn模块实现了网络层类,包括卷积层(Conv2d)、最大池化层(MaxPool2d)、全连接层(Linear)和其他激活层

torch.nn模块提供functional类用于网络层类的实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class LeNet5(nn.Module):

def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0, bias=True)
self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0, bias=True)
self.conv3 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5, stride=1, padding=0, bias=True)

self.pool = nn.MaxPool2d((2, 2), stride=2)

self.fc1 = nn.Linear(in_features=120, out_features=84, bias=True)
self.fc2 = nn.Linear(84, 10, bias=True)

def forward(self, input):
x = self.pool(F.relu(self.conv1(input)))
x = self.pool(F.relu(self.conv2(x)))
x = self.conv3(x)

x = x.view(-1, self.num_flat_features(x))

x = F.relu(self.fc1(x))
return self.fc2(x)

def num_flat_features(self, x):
size = x.size()[1:] # all dimensions except the batch dimension
num_features = 1
for s in size:
num_features *= s
return num_features

训练

训练参数如下

  1. 学习率lr = 1e-3
  2. 批量大小batch_size = 128
  3. 迭代次数epochs = 500

训练结果

训练时间

CPUGPU单次迭代时间
8核 Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHzGeForce 940MX约13秒

迭代500次训练结果

训练集精度测试集精度
99.40%98.63%

相关阅读