本文代码完全借鉴pytorch中文手册
'''我们找到数据集,对数据做预处理,定义我们的模型,调整超参数,测试训练,再通过训练结果对超参数进行调整或者对模型进行调整。'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim #实现各种优化算法的库
from torchvision import datasets,transformsBATCH_SIZE=512 #大概需要2G的显存
EPOCHS=20 #总共训练20次
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu") #让torch判断是否使用GPU#对数据进行预处理
transforms=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))])
#准备数据集,路径是关于py文件的相对路径
trainset=datasets.MNIST(root='./MNIST_data',train=True,download=False,transform=transforms)#加载数据集
train_loader=torch.utils.data.DataLoader(trainset,batch_size=BATCH_SIZE,shuffle=True)#准备测试集
testset=datasets.MNIST(root='./MNIST_data',train=True,download=False,transform=transforms)#加载测试集
test_loader=torch.utils.data.DataLoader(testset,batch_size=BATCH_SIZE,shuffle=True)#定义卷积神经网络
class ConvNet(nn.Module):def __init__(self):super().__init__()#batch*1*28*28(每次会送入batch个样本,输入通道数1(黑白图像)),图像分辨率是28*28)#下面的卷积层Conv2d的第一个参数指输入通道数,第二个参数指输出通道数,第三个参数指卷积核的大小self.conv1=nn.Conv2d(1,10,5)self.conv2=nn.Conv2d(10,20,3)#下面的全连接层Linear的第一个参数指输入通道数,第二个参数指输出通道数self.fc1=nn.Linear(20*10*10,500) #输入通道数是2000,输出通道数是500self.fc2=nn.Linear(500,10) #输入通道数是500,输出通道数是10,即10分类def forward(self,x):in_size=x.size(0) #在本例中in_size=512,也就是BATCH_SIZE的值。输入的x可以看成是512*1*28*28的张量out=self.conv1(x) #batch*1*28*28 -> batch*10*24*24(28×28的图像经过一次核为5×5的卷积,输出变为24×24)out=F.relu(out) #batch*10*24*24(激活函数ReLU不改变形状)out=F.max_pool2d(out,2,2)#batch*10*24*24 -> batch*10*12*12(2×2的池化层会减半)out=self.conv2(out) #batch*10*12*12 -> batch*20*10*10(再卷积一次,核的大小是3)out=F.relu(out)out=out.view(in_size,-1) #batch*20*10*10 -> batch*2000(out的第二维是-1,说明是自动推算,本例中第二维是20*10*10)out=self.fc1(out) #batch*2000 -> batch*500out=F.relu(out) out=self.fc2(out) #batch*500 -> batch*10out=F.log_softmax(out,dim=1) #计算log(softmax(x)),用log是为了防止数过大。return out#我们实例化一个网络,实例化后使用.to方法将网络移动到GPU
model=ConvNet().to(DEVICE)#优化器我们也直接选择简单暴力的Adam
optimizer=optim.Adam(model.parameters())#定义一个训练函数
def train(model,device,train_loader,optimizer,epoch):model.train() #启用BatchNormalization和Dropout,将BatchNormalization和Dropout置为Truefor batch_idex,(data,target) in enumerate(train_loader): ##将迭代器的数据组成一个索引系列,并输出索引和值,batch_diex是序号,后者是数据data,target=data.to(device),target.to(device) #在gpu上跑optimizer.zero_grad() #梯度清零output=model(data) #将数据放入模型loss=F.nll_loss(output,target) #计算损失函数loss.backward() #计算梯度optimizer.step()if (batch_idex+1)%30 == 0: #每训练30个打印一次print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(epoch,batch_idex*len(data),len(train_loader.dataset),100. * batch_idex/len(train_loader.dataset),loss.item()))#定义一个测试函数
def test(model,device,test_loader):model.eval() #不启用BatchNormalization和Dropout,将BatchNormalization和Dropout置为Falsetest_loss=0correct=0with torch.no_grad():for data,target in test_loader:data,target=data.to(device),target.to(device) #在gpu上跑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)