我学Pytorch之五~~~PyTorch实现Minist（CNN版）

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())