今天参考了-深度学习框架PyTorch:入门与实践,这本书自己也动手写写LeNet。
LeNet 网络结构如下:
PyTorch 使用 torchvision 自带的 CIFAR10 数据实现。
运行环境:pytorch 0.4.0 CPU版、Python 3.6、Windows 7
# Pytorch 0.4.0 LeNet实现cifar10分类
# @Time :2018/6/14
# @Author :LiuYinxing
import torchvision as tv
import torchvision.transforms as transforms
import torch as t
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
t.set_num_threads(8)
class Net(nn.Module): # 网络结构
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = x.view(x.size()[0], -1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
def getData(): # 定义对数据的预处理
transform = transforms.Compose([
transforms.ToTensor(), # 转为Tensor
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), # 归一化
])
trainset = tv.datasets.CIFAR10(root='/data/', train=True, download=True, transform=transform) # 训练集
trainloader = t.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)
testset = tv.datasets.CIFAR10('/data/', train=False, download=True, transform=transform) # 测试集
testloader = t.utils.data.DataLoader(testset, batch_size=4, shuffle=False)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
return trainloader, testloader, classes
def trainModel(): # 训练模型
trainloader, testloader, _ = getData() # 获取数据
net = Net()
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # 定义优化器
for epoch in range(2):
for step, (tx, ty) in enumerate(trainloader, 0):
optimizer.zero_grad() # 梯度清零
py = net(tx) # forward + backward
loss = criterion(py, ty)
loss.backward()
optimizer.step() # 更新参数
if step % 1000 == 999:
acc = testNet(net, testloader)
print('Epoch:', epoch, '|Step:', step, '|train loss:%.4f' % loss.item(), '|test accuracy:%.4f' % acc)
print('Finished Training')
return net
def testNet(net, testloader): # 获取在测试集上的准确率
correct, total = .0, .0
for x, y in testloader:
py = net(x)
_, predicted = t.max(py, 1) # 获取分类结果
total += y.size(0) # 记录总个数
correct += (predicted == y).sum() # 记录分类正确的个数
return float(correct) / total
if __name__ == '__main__':
net = trainModel()
注:
testloader = t.utils.data.DataLoader(
testset,
batch_size=4,
shuffle=False,
num_workers=2)
num_workers=2 # 多线程来读数据,在Windows使用时应该特别注意,很容易出错。
欢迎指正哦。