本文记录了pytorch训练MNIST数据集的过程,通过本文可熟悉pytorch训练的大体操作过程。
一、导入各种模块
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
二、设置超参数
BATCH_SIZE=512 # 批次大小
EPOCHS=20 # 总共训练批次
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 让torch判断是否使用GPU,建议使用GPU环境,因为会快很多
二、下载并读取数据集
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=BATCH_SIZE, shuffle=True)
torch.utils.data.DataLoader是PyTorch中数据读取的一个重要接口,只要是用PyTorch来训练模型基本都会用到该接口,该接口主要用来将自定义的数据读取接口的输出或者PyTorch已有的数据读取接口的输入按照batch size封装成Tensor。
torch.utils.data.DataLoader的重要参数:
1、dataset,这个是PyTorch已有的数据读取接口(比如torchvision.datasets.ImageFolder(可用于自己的图片数据集的训练))或者自定义的数据接口的输出,该接口是torch.utils.data.Dataset类的对象,或者是继承自torch.utils.data.Dataset类的自定义类的对象。
2、batch_size,训练批次的大小,根据具体情况设置即可。
3、shuffle,在封装过程中是否打乱数据集,一般在训练数据中会采用。
具体见此博客(注:现在版本不需要包装成Variable)
三、定义卷积神经网络
class ConvNet(nn.Module):
def __init__(self):
super().__init__()
# 1,28x28
self.conv1=nn.Conv2d(1 ,10,5) # 24x24
self.pool = nn.MaxPool2d(2,2) # 12x12
self.conv2=nn.Conv2d(10,20,3) # 10x10
self.fc1 = nn.Linear(20*10*10,500)
self.fc2 = nn.Linear(500,10)
def forward(self,x):
in_size = x.size(0)
out = self.conv1(x) #24
out = F.relu(out)
out = self.pool(out) #12
out = self.conv2(out) #10
out = F.relu(out)
out = out.view(in_size,-1)
out = self.fc1(out)
out = F.relu(out)
out = self.fc2(out)
out = F.log_softmax(out,dim=1)
return out
关于定义网络这一部分,详细介绍可以查看我的另一篇博客pytorch中根据神经网络结构确定输入图片尺寸(根据图片尺寸修改神经网络结构)
四、神经网络的训练与评估
model = ConvNet().to(DEVICE)
optimizer = optim.Adam(model.parameters())
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)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if(batch_idx+1)%30 == 0:
print('Train Epoch: {} [{}/{} ({:.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() # 将一批的损失相加
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)