06、使用PyTorch在Mnist数据集上训练与验证的所有流程实现

东方耀 · 发表于 2020-4-13 18:40:48

# -*- coding: utf-8 -*-
__author__ = u'东方耀微信：dfy_88888'
__date__ = '2020/4/13 下午5:11'
__product__ = 'PyCharm'
__filename__ = 'dfy03_mnist'
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
print(torch.__version__)
BATCH_SIZE = 512 # 大概需要2G的显存
EPOCHS = 10 # 总共训练批次
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 让torch判断是否使用GPU，建议使用GPU环境，因为会快很多
# 对图片数据进行预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.1307,), std=(0.3081,))
])
trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
print(type(trainset))
# pytorch中的数据加载器DataLoader
train_loader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE, shuffle=True)
print(type(train_loader))
minist01 = iter(train_loader) # 迭代器方便按批次取数据
imgs, labels = minist01.next()
print(type(imgs), imgs.numpy().shape, imgs.size())
print(type(labels), labels.size())
class ConvNet(nn.Module):
def __init__(self):
super().__init__()
# 1,28x28
self.conv1 = nn.Conv2d(1, 10, 5) # 10, 24x24
self.conv2 = nn.Conv2d(10, 20, 3) # 20, 10x10
self.fc1 = nn.Linear(20 * 10 * 10, 500)
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
in_size = x.size(0)
# in_size = batch_size
out = self.conv1(x) # 24
out = F.relu(out)
out = F.max_pool2d(out, 2, 2) # 12
out = self.conv2(out) # 10
out = F.relu(out)
# reshape flatten
out = out.view(in_size, -1) # [None, 2000]
out = self.fc1(out)
out = F.relu(out)
out = self.fc2(out)
out = F.log_softmax(out, dim=1)
return out
model = ConvNet().to(DEVICE)
for i, (name, param) in enumerate(model.named_parameters()):
print('名字是:%s, 第%d个的权重参数的size=' % (name, i), param.size())
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 下面定义一下训练的函数，我们将训练的所有操作都封装到这个函数中
def train(model, device, train_loader, optimizer, epoch):
# 设置模型的训练模式
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
# 梯度归0
optimizer.zero_grad()
# 模型前向计算
output = model(data)
# 这是什么loss？ The negative log likelihood loss
# loss = F.nll_loss(output, target)
loss = F.cross_entropy(output, target)
loss.backward() # 反向传播算梯度
optimizer.step() # 更新权重参数
if (batch_idx + 1) % 30 == 0:
print('第{}轮Train [{}/共{}个训练样本 ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
# 测试的操作也一样封装成一个函数
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
# test_loss += F.nll_loss(output, target, reduction='sum').item() # 将一批的损失相加
test_loss += F.cross_entropy(output, target, reduction='sum').item()
pred = output.max(1, keepdim=True)[1] # 找到概率最大的下标
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/共{}个验证样本 ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
# 下面开始训练，这里就体现出封装起来的好处了，只要写两行就可以了
for epoch in range(1, EPOCHS + 1):
train(model, DEVICE, train_loader, optimizer, epoch)
test(model, DEVICE, test_loader)
# MNIST是一个很简单的数据集，由于它的局限性只能作为研究用途，对实际应用带来的价值非常有限。
# 但是通过这个例子，我们可以完全了解一个实际项目的工作流程
#
# 我们找到数据集，对数据做预处理，定义我们的模型，调整超参数，测试训练，
# 再通过训练结果对超参数进行调整或者对模型进行调整。
#
# 并且通过这个实战我们已经有了一个很好的模板，以后的项目都可以以这个模板为样例

复制代码

		自动登录	找回密码
密码			立即注册

[PyTorch] 06、使用PyTorch在Mnist数据集上训练与验证的所有流程实现