Pytorch如何训练MNIST模型

1、在python环境中先导入Pytorch相关的包, 在Conda环境下可使用 conda install pytorch torchvision cudatoolkit=10.1 -c pytorch 命令安装Pytorch

2、使用Pytorch导入MNIST数据集, 第一次导入时设置download=True可以自动下载MNIST数据集,transform将MNIST数据集归一化,batchsize设置为128

3、定义一个名为PytorchMnistModel的网络模型, 输入为28*28维, 输出为10维, 包含fc1和fc2两个全连接层, conv1和conv2两个卷积层

4、实例化上一步定义的PytorchMnistModel模型,.cuda()将模型放到GPU上, criterion损失函数设置为交叉熵CrossEntropyLoss,optimizer优化器设置为Adam学习率lr设置为0.001

5、range(15)将epoch迭代次数设置为15,.cuda()将数据放到GPU上,training_loss记录训练损失的数值, plt.plot画出损失大小的折线图,display每训练一个batch就将折线图刷新一次