Mnist手写数据集的模型搭建及GPU训练(基于pytorch)

1. 用意

pytorch看了有一段时间了，可是怎么熟练操作pytorch呢，于是就用手写数据集来办事，键盘滚起来！
四颗GPU训练
两个类完成训练

代码直接copy过去，立即跑通，但电脑上必须有GPU
代码直接copy过去，立即跑通，但电脑上必须有GPU
代码直接copy过去，立即跑通，但电脑上必须有GPU

2.结果展示

2.1 四颗GPU

2.2 训练过程

2.3 训练结果

3.代码如下

import os
import random
import torch
import numpy as np
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(  # 输入图片 shape (1, 28, 28)
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2,
            ),  # 卷积后的shape (16, 28, 28)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),   # 池化之后的shape (16, 14, 14)
        )
        self.conv2 = nn.Sequential(  # input shape (16, 14, 14)
            nn.Conv2d(16, 32, 5, 1, 2),  # output shape (32, 14, 14)
            nn.ReLU(),
            nn.MaxPool2d(2),  # output shape (32, 7, 7)
        )
        self.out = nn.Linear(32 * 7 * 7, 10)  # 全连接层，输出为10

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # flatten the output of conv2 to (batch_size, 32 * 7 * 7)
        output = self.out(x)
        return output


class TrainMnistDataset:
    def __init__(self, epoch=1000, lr=0.001,  batch_size=50, data_root_dir="./datasets/", model_save_dir="./model_save_dir",
                 pretraind_model_path=False):
        self.data_root_dir = data_root_dir
        self.batch_size