【Pytorch】MNIST数据集的训练和测试

训练和测试的完整代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader, Dataset
import argparse
import os# 训练
def train(args, model, device, train_loader, optimizer):model.train()num_correct = 0for batch_index, (images, labels) in enumerate(train_loader):images = images.to(device)labels = labels.to(device)# forwardoutputs = model(images)loss = F.cross_entropy(outputs, labels)# backwardoptimizer.zero_grad()   # 梯度清空loss.backward()            # 梯度回传，更新参数optimizer.step()_, predicted = torch.max(outputs, dim=1)# 每一个batch预测对的个数batch_correct = (predicted == labels).sum().item()# 每一个batch的准确率batch_accuracy = batch_correct / args.batch_size# 每一个epoch预测对的总个数num_correct += (predicted == labels).sum().item()# print sth.print(f'Epoch:{epoch},Batch ID:{batch_index}/{len(train_loader)}, loss:{loss}, Batch accuracy:{batch_accuracy*100}%')# 每一个epoch的准确率epoch_accuracy = num_correct / len(train_loader.dataset)# print epoch_accuracyprint(f'Epoch Accuracy:{epoch_accuracy}')# 保存模型if epoch % args.checkpoint_interval == 0:torch.save(model.state_dict(), f"checkpoints/VGG16_MNIST_%d.pth" % epoch)# 验证
def test(args, model, device, test_loader):model.eval()total_loss = 0num_correct = 0if args.pretrained_weights.endswith(".pth"):model.load_state_dict(torch.load(args.pretrained_weights))# 不计算梯度，节省计算资源with torch.no_grad():for images, labels in test_loader:images = images.to(device)labels = labels.to(device)output = model(images)# 总的losstotal_loss += F.cross_entropy(output, labels).item()     # item()用于取出tensor里边的值# torch.max():返回的是两个值，第一个值是具体的value，第二个值是value所在的index_, predicted = torch.max(output, dim=1)# 预测对的总个数num_correct += (predicted == labels).sum().item()# 平均losstest_loss = total_loss / len(test_loader.dataset)# 平均准确率accuracy = num_correct / len(test_loader.dataset)# print sth.print(f'Average loss:{test_loss}\nTest Accuracy:{accuracy*100}%')if __name__ == '__main__':parser = argparse.ArgumentParser(description = 'Pytorch-MNIST_classification')parser.add_argument('--epochs', type=int, default=20, help='number of epochs')parser.add_argument('--batch_size', type=int, default=32, help='size of each image batch' )parser.add_argument('--num_classes', type=int, default=10, help='number of classes')parser.add_argument('--lr', type=float, default=0.001, help='learning rate')parser.add_argument('--momentum', type=float, default=0.9, help='SGD momentum')parser.add_argument('--pretrained_weights', type=str, default='checkpoints/', help='if specified starts from checkpoint model')parser.add_argument("--img_size", type=int, default=224, help="size of each image dimension")parser.add_argument("--checkpoint_interval", type=int, default=1, help="interval between saving model weights")parser.add_argument("--train",  default=False, help="train or test")args = parser.parse_args()print(args)device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# os.makedirs() 方法用于递归创建目录os.makedirs("output", exist_ok=True)os.makedirs("checkpoints", exist_ok=True)# transformdata_transform = transforms.Compose([transforms.ToTensor(), transforms.RandomResizedCrop(args.img_size)])# 下载训练数据train_data = datasets.MNIST(root = 'data',train = True,transform = data_transform,target_transform = None,download = True)# 下载测试数据test_data = datasets.MNIST(root = 'data',train = False,transform = data_transform,target_transform = None,download = True)# 加载训练数据train_loader = DataLoader(dataset = train_data,batch_size = args.batch_size,shuffle = True)# 加载测试数据test_loader = DataLoader(dataset = test_data,batch_size = args.batch_size)# 创建模型model = models.vgg16(pretrained = True)# 修改vgg16的输出维度model.classifier[6] = nn.Linear(in_features=4096, out_features=args.num_classes, bias=True)# MNIST数据集是灰度图，channel数为1model.features[0] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))print(model)model = model.to(device)# 优化器(也可以选择其他优化器)optimizer = torch.optim.SGD(model.parameters(), lr = args.lr, momentum = args.momentum)# optimizer = torch.optim.Adam()if args.train == True:for epoch in range(1, args.epochs+1):# 是否加载预训练好的权重if args.pretrained_weights.endswith(".pth"):model.load_state_dict(torch.load(args.pretrained_weights))train(args, model, device, train_loader, optimizer)else:# 是否加载预训练好的权重if args.pretrained_weights.endswith(".pth"):model.load_state_dict(torch.load(args.pretrained_weights))test(args, model, device, test_loader)

测试结果：
我只训练了不到10轮，效果不是太好，还有提升空间。

说明：
MNIST数据集可以通过trochvision中的datasets.MNIST下载，也可以自己下载（注意存放路径）；我模型使用的是torchvision中的models中预训练好的vgg16网络，也可以自己搭建网络。

【Pytorch】MNIST数据集的训练和测试相关推荐

【caffe】mnist数据集lenet训练与测试
在上一篇中,费了九牛二虎之力总算是把Caffe编译通过了,现在我们可以借助mnist数据集,测试下Caffe的训练和检测效果. 准备工作:在自己的工作目录下,新建一个文件夹,命名为mnist_test ...
【Pytorch分布式训练】在MNIST数据集上训练一个简单CNN网络，将其改成分布式训练
文章目录普通单卡训练-GPU 普通单卡训练-CPU 分布式训练-GPU 分布式训练-CPU 租GPU服务器相关以下代码示例基于:在MNIST数据集上训练一个简单CNN网络,将其改成分布式训练. 普 ...
在MNIST数据集上训练一个手写数字识别模型
使用Pytorch在MNIST数据集上训练一个手写数字识别模型, 代码和参数文件可下载 1.1 数据下载 import torchvision as tvtraining_sets = tv.dat ...
DL之DNN：自定义2层神经网络TwoLayerNet模型(计算梯度两种方法)利用MNIST数据集进行训练、预测
DL之DNN:自定义2层神经网络TwoLayerNet模型(计算梯度两种方法)利用MNIST数据集进行训练.预测导读利用python的numpy计算库,进行自定义搭建2层神经网络TwoLayerN ...
FCN制作自己的数据集、训练和测试 caffe
原文:http://blog.csdn.net/zoro_lov3/article/details/74550735 FCN制作自己的数据集.训练和测试全流程花了两三周的时间,在导师的催促下,把FC ...
DL之DNN：自定义2层神经网络TwoLayerNet模型(封装为层级结构)利用MNIST数据集进行训练、预测
DL之DNN:自定义2层神经网络TwoLayerNet模型(封装为层级结构)利用MNIST数据集进行训练.预测导读计算图在神经网络算法中的作用.计算图的节点是由局部计算构成的. ...
DL之DNN：自定义2层神经网络TwoLayerNet模型(封装为层级结构)利用MNIST数据集进行训练、GC对比
DL之DNN:自定义2层神经网络TwoLayerNet模型(封装为层级结构)利用MNIST数据集进行训练.GC对比导读神经网络算法封装为层级结构的作用.在神经网络算法中,通过将 ...
DL之DCGNN：基于TF利用DCGAN实现在MNIST数据集上训练生成新样本
DL之DCGNN:基于TF利用DCGAN实现在MNIST数据集上训练生成新样本目录输出结果设计思路实现部分代码说明:所有图片文件丢失输出结果更新-- 设计思路更新-- 实现部分代码更 ...
python划分数据集用pandas_用pandas划分数据集实现训练集和测试集
1.使用model_select子模块中的train_test_split函数进行划分数据:使用kaggle上Titanic数据集划分方法:随机划分 # 导入pandas模块,sklearn中mo ...

【Pytorch】MNIST数据集的训练和测试

【Pytorch】MNIST数据集的训练和测试相关推荐

最新文章

热门文章