整个过程如下:
- 模型定义
- 数据处理和加载
- 训练模型
- 训练过程的可视化
- 测试
1.万事先导包
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import mnist
from torch.autograd import Variable
2.模型定义
定义一个LeNet模型,模型详解参见:LeNet详解
卷积->池化->卷积->池化->全连接->全连接->全连接
#定义LeNet网络模型
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1,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(-1,16*5*5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
3.数据处理和加载
数据处理:先对传入的MNIST数据集中的图片进行transform处理
LeNet默认输入32×32大小的图片,而MNIST提供的图片是28*28,故需要Resize(32)
#定义训练集图片格式
#注:LeNet默认输入32×32大小的图片,而MNIST提供的图片是28*28,故需要Resize(32)
transform = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(), #image转为Tensor
transforms.Normalize([0.5], [0.5])]) #归一化
数据加载,定义训练集和测试集的dataset、dataloader
- datasets.MNIST()的第一个参数为MNIST数据集的存放路径,
第二个参数train=True为训练集,train=False为测试集,
第三个参数transform为对数据集图片的处理。 - DataLoader的第一个参数为传入数据集,
第二个参数shuffle=True为打乱数据集内数据顺序,
第三个参数batch_size为一次训练几张图片,
第四个参数num_workers为多线程任务数。
#定义训练集,从MNIST中自动获取
trainset = datasets.MNIST('resourses/',train=True,transform=transform,download=False)
trainloader = DataLoader(trainset,shuffle=True,batch_size=5,num_workers=0)
#定义测试集,从MNIST中自动获取
testset = datasets.MNIST('resourses/',train=False,transform=transform)
testloader = DataLoader(testset,shuffle=True,batch_size=5,num_workers=0)
初始化模型,定义损失函数和优化器
USE_CUDA = torch.cuda.is_available()
lenet = LeNet()
#定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(lenet.parameters(),lr=0.01)
4.训练模型
训练模型步骤:
①:将数据从DataLoader取出用于迭代,注意,数据要转换成Variable,用于可求导运算
②:如果cuda可用,将模型及数据放到cuda上跑
③:梯度清零(放在哪个位置都可以,只要在下次计算梯度之前即可)
④:forward + backward
将从训练集取出的图片传进模型,获取预测标签
定义loss function
反向传播
⑤:更新参数
#训练
print("Begin Training")
for epoch in range(4):
running_loss = 0.0
running_acc = 0.0
#训练
for i,data in enumerate(trainloader,0):
images,labels = data
images,labels = Variable(images),Variable(labels)
if USE_CUDA:
#如果cuda可用,将模型及数据搬到cuda上计算
lenet = lenet.cuda()
images,labels = images.cuda(),labels.cuda()
#梯度清零
optimizer.zero_grad()
#forward + backward
pre_labels = lenet(images)
loss = criterion(pre_labels,Variable(labels))
loss.backward()
#更新参数
optimizer.step()
5.训练过程中的可视化:
接上面代码
格式化输出训练loss及训练正确率accuracy
#打印log信息
running_loss += loss.item()
_,pred = torch.max(pre_labels,1)
num_correct = (pred == labels).sum()
accuracy = (pred == labels).float().mean()
running_acc += num_correct.item()
print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(
epoch + 1, running_loss / (len(trainset)), running_acc / (len(trainset))))
6.测试模型
接上面代码
格式化输出测试loss及测试正确率accuracy
#测试
lenet.eval()
eval_loss = 0
eval_acc = 0
#测试
for data in testloader:
img, label = data
if USE_CUDA:
lenet = lenet.cuda()
img, label = img.cuda(), label.cuda()
out = lenet(img)
loss = criterion(out, label)
eval_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
eval_acc += num_correct.item()
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(testset)), eval_acc / (len(testset))))
print('Finished Training and Testing')
完整代码附上
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import mnist
from torch.autograd import Variable
#定义LeNet网络模型
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1,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(-1,16*5*5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
#定义训练集图片格式
#注:LeNet默认输入32×32大小的图片,而MNIST提供的图片是28*28,故需要Resize(32)
transform = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(), #image转为Tensor
transforms.Normalize([0.5], [0.5])]) #归一化
#定义训练集,从MNIST中自动获取
trainset = datasets.MNIST('resourses/',train=True,transform=transform,download=False)
trainloader = DataLoader(trainset,shuffle=True,batch_size=5,num_workers=0)
#定义测试集,从MNIST中自动获取
testset = datasets.MNIST('resourses/',train=False,transform=transform)
testloader = DataLoader(testset,shuffle=True,batch_size=5,num_workers=0)
USE_CUDA = torch.cuda.is_available()
lenet = LeNet()
#定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(lenet.parameters(),lr=0.01)
#训练+测试
print("Begin Training")
for epoch in range(4):
running_loss = 0.0
running_acc = 0.0
#训练
for i,data in enumerate(trainloader,0):
images,labels = data
images,labels = Variable(images),Variable(labels)
if USE_CUDA:
#如果cuda可用,将模型及数据搬到cuda上计算
lenet = lenet.cuda()
images,labels = images.cuda(),labels.cuda()
#梯度清零
optimizer.zero_grad()
#forward + backward
pre_labels = lenet(images)
loss = criterion(pre_labels,Variable(labels))
loss.backward()
#更新参数
optimizer.step()
#打印log信息
running_loss += loss.item()
_,pred = torch.max(pre_labels,1)
num_correct = (pred == labels).sum()
accuracy = (pred == labels).float().mean()
running_acc += num_correct.item()
print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(
epoch + 1, running_loss / (len(trainset)), running_acc / (len(trainset))))
lenet.eval()
eval_loss = 0
eval_acc = 0
#测试
for data in testloader:
img, label = data
if USE_CUDA:
lenet = lenet.cuda()
img, label = img.cuda(), label.cuda()
out = lenet(img)
loss = criterion(out, label)
eval_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
eval_acc += num_correct.item()
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(testset)), eval_acc / (len(testset))))
print('Finished Training and Testing')