import torch
from torch import nn,optim
import random
import numpy as np
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import os
import torch.utils.data as Data
#1、准备数据
batch_size=50
input_size=784
num_classes=10
num_epochs=5
learnning_rate=0.001
hidden_size=200
layer1_out_channels=2
layer2_out_channels=4
layer3_out_channels=16
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_data=torchvision.datasets.MNIST(root='../../data',train=True,transform=transforms.ToTensor(),download=True)
train_loader=Data.DataLoader(dataset=train_data,batch_size=batch_size,shuffle=True)
test_data=torchvision.datasets.MNIST(root='../../data',train=False,transform=transforms.ToTensor(),download=True)
test_loader=Data.DataLoader(dataset=test_data,batch_size=len(test_data))
#2、建立模型
class ConvNet(nn.Module):
def __init__(self,num_classes):
super(ConvNet,self).__init__()
#第一层隐藏层
self.layer1=nn.Sequential(
nn.Conv2d(1,layer1_out_channels,kernel_size=3,stride=1,padding=1), #卷积,满足公式计算:其中n为原始图像大小(宽高对应),p为Padding填充维度,f伪卷积核维度,s为步长。。。(n+2p-f)/s
nn.BatchNorm2d(layer1_out_channels), #标准化
nn.ReLU(), #激活函数
nn.MaxPool2d(kernel_size=2,stride=2) #最大池化函数(抽样)
)
#第二层隐藏层
self.layer2=nn.Sequential(
nn.Conv2d(layer1_out_channels,layer2_out_channels,kernel_size=5,stride=1,padding=2),
nn.BatchNorm2d(layer2_out_channels),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2,stride=2)
)
#第三层隐藏层
self.layer3=nn.Sequential(
nn.Conv2d(layer2_out_channels,layer3_out_channels,kernel_size=3,stride=1,padding=2),
nn.BatchNorm2d(layer3_out_channels),
nn.ReLU(),
nn.MaxPool2d(kernel_size=1,stride=2)
)
#全连接层
self.fc=nn.Sequential(
nn.Linear(5*5*layer3_out_channels,num_classes)
)
def forward(self,x):
x=self.layer1(x)
out=self.layer2(x)
out=self.layer3(out)
out=torch.reshape(out,shape=(out.shape[0],-1))
out=self.fc(out)
return out
#3、模型训练
isloadmodel=0
criterion=nn.CrossEntropyLoss()
if isloadmodel==0:
model=ConvNet(num_classes)
optimizer=optim.Adam(model.parameters(),lr=learnning_rate)
total_steps=len(train_loader)
loss_arr=[]
for epoch in range(num_epochs):
for i,(images,labels) in enumerate(train_loader):
images=images.to(device)
labels=labels.to(device)
outputs=model(images)
loss=criterion(outputs,labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1)%batch_size==0:
plt.ion()
print("epoch:{}/{},steps:{}/{},loss:{}".format(epoch+1,num_epochs,i+1,total_steps,loss.item()))
loss_arr.append(loss.item())
#损失可视化
plt.cla()
plt.plot(loss_arr)
plt.text(len(loss_arr)/2,np.max(loss_arr)/2,"loss:{}".format(round(loss.item(),2)))
plt.pause(0.2)
plt.ioff()
plt.show()
torch.save(model,'model.ckpt')
else:
model=torch.load('model.ckpt')
#4、测试模型
with torch.no_grad():
for images,labels in test_loader:
# images=torch.reshape(images,shape=(-1,input_size))
images=images.to(device)
labels=labels.to(device)
outputs=model(images)
loss=criterion(outputs,labels)
_,predit=torch.max(outputs,1)
accuracy=(predit==labels).sum()/float(len(test_data))
print('Accuracy:{},loss:{}',accuracy.item(),loss.item())