深度学习基础实战使用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  #大概需要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,输出通道数是500
		self.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*500
		out=F.relu(out)         
		out=self.fc2(out)		 #batch*500 -> batch*10
		out=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置为True
	for 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置为False
	test_loss=0
	correct=0
	with 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)