卷积神经网络pytorch实现
一、模型建立
本代码基于pytorch以及cuda的GPU模式,如果没有GPU的话,把to.devide()这种代码删去就好。
下面仅仅是代码的实现,不对模型建立做解释,相关的代码解释有注释,希望对大家有帮助。
1 Lenet5卷积神经网络模型
需要先运行模型,然后下面调用才能成功。其次,模型代码文件需与运行的代码文件在同一级同一目录下。
from torch import nn
class Lenet5(nn.Module):
"""
for cifar10 dataset
"""
def __init__(self):
super(Lenet5, self).__init__()
self.conv_unit = nn.Sequential(
nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
# 输入是(batch size,3,32,32), 上面3是输入通道数,6是输出通道数
# 经过第一层之后,输出是(batch size, 6, 28, 28)
nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
# 这个就是调用max pooling
# 此时得输出是(batch size, 6, 14, 14)
nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
# 输出(batch size, 16, 5, 5)
)
# flatten
# fc unit
self.fc_unit = nn.Sequential(
nn.Linear(16*5*5, 120), # 不用乘batch size
nn.ReLU(),
nn.Linear(120, 84),
nn.ReLU(),
nn.Linear(84, 10)
)
self.criteon = nn.CrossEntropyLoss() # 评价函数
def forward(self, x):
"""
:param x: [batch size, 3, 32, 32]
:return:
"""
batchsz = x.size(0)
x = self.conv_unit(x)
# 把x从[batch size, 16, 5, 5],变成[batch size, 16*5*5]
x = x