您的位置:首页 > 新闻 > 会展 > minist数据集分类模型的训练

minist数据集分类模型的训练

2025/3/10 17:31:58 来源:https://blog.csdn.net/Laollaoaolao/article/details/140245626  浏览:    关键词:minist数据集分类模型的训练

minist数据集训练

训练方法:利用pytorch来实现minist数据集的分类模型训练

训练模型如下图所示

img

模型代码:

import torch
from torch import nn
from torch.nn import Flattenclass Net(nn.Module):def __init__(self):super().__init__()self.module = nn.Sequential(nn.Conv2d(1,4,5,1,0),nn.MaxPool2d(2),nn.Conv2d(4,8,5,1,0),nn.MaxPool2d(2),# torch.Size([8, 4, 4])Flatten(),# torch.Size([64, 128])nn.Linear(8*16,16),# torch.Size([64, 16])nn.Linear(16, 10) #torch.Size([64, 10]))def forward(self, x):x = self.module(x)return xif '__name__' == '__mian__':net = Net()input = torch.ones((64,1, 28, 28))output = net(input)print(output.shape)

现在开始编写代码训练模型

基本思路:

  1. 读取数据集并整理
  2. 将数据集放入模型中训练,每次记录训练的损失loss以及准确率

实测训练大概10轮后,正确率能达到98%以上

import torch.optim
import  torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import  SummaryWriter
from module import *
dataset_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])train_set = torchvision.datasets.MNIST(root='../dataset_MNIST',train=True,transform=dataset_transform,download=True)
test_set = torchvision.datasets.MNIST(root='../dataset_MNIST',train=False,transform=dataset_transform,download=True)print('训练数据集size{}'.format(len(train_set)))
print('测试数据集size{}'.format(len(test_set)))train_loader = DataLoader(dataset=train_set, batch_size=64)
test_loader = DataLoader(dataset=test_set,batch_size=64)# for data in train_loader:
#     imgs,target = data
#     print('imgs',imgs.shape) # torch.Size([64, 1, 28, 28])
#     print('target',target)
#     breakwriter = SummaryWriter('../p10')
cuda_available = torch.cuda.is_available()net = Net()
if cuda_available:net = net.cuda()# 定义损失函数
loss_fn = nn.CrossEntropyLoss()
if cuda_available:loss_fn = loss_fn.cuda()# 优化器
learning_rate = 0.01
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)# 设置训练参数
# 训练次数
total_train_step = 0
total_test_step = 0
# 训练轮数
epoch = 10for i in range(10):print('------------第{}轮训练开始------------'.format(i))net.train()for data in train_loader:imgs,targets = dataif cuda_available:imgs = imgs.cuda()targets = targets.cuda()outputs = net(imgs)loss = loss_fn(outputs,targets)#优化器优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_train_step += 1if total_train_step % 100 == 0:print('训练次数:{} loss:{}'.format(total_train_step, loss.item()))# 开始测试net.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_loader:imgs,targets = dataif cuda_available:imgs = imgs.cuda()targets = targets.cuda()outputs = net(imgs)loss = loss_fn(outputs, targets)total_test_loss += loss.item()accuracy = (outputs.argmax(1) == targets).sum()total_accuracy += accuracyprint("整体测试集的Loss:{}".format(total_test_loss))writer.add_scalar("test_loss", total_test_loss, total_test_step)total_test_step += 1print("整体测试集的正确率:{}".format(total_accuracy / len(test_set)))if i == 9:torch.save(net, "train_model_{}.pth".format(i))print('模型已保存')writer.close()

上面已经将最后一次训练的的模型保存了,那么现在就可以用一个图片来测试一下这个模型

from module import *
import torchvision.transforms
from PIL import  Image
dataset_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
test_set = torchvision.datasets.MNIST('../dataset_MNIST',train=False,transform=dataset_transform,download=False)model = torch.load("train_model_9.pth")
model.eval()
for i in range(10):img, target = test_set[i]print('--------------第{}张图片--------------'.format(i))print("图片尺寸",img.shape)print("标签",target)img = torch.reshape(img,(1,1,28,28))img = img.cuda()with torch.no_grad():output = model(img)print(output)print(output.argmax(1))

这里要注意的一点是,因为训练模型时使用了nvida的cuda驱动,那么在使用模型的时候,也需要将图片进行一点修改,具体是

 img = img.cuda()

测试结果如下:

--------------第0张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 7
tensor([[  0.1551,  -2.2232,   5.1375,   3.1896,  -9.0812,  -3.7413, -16.3016,16.5437,  -2.3190,   2.7608]], device='cuda:0')
tensor([7], device='cuda:0')
--------------第1张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 2
tensor([[  8.5477,   6.7017,  15.3023,   6.9958, -12.9011,  -3.1987,   2.8130,-12.9218,   4.8305, -17.2936]], device='cuda:0')
tensor([2], device='cuda:0')
--------------第2张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 1
tensor([[-1.4329,  9.1416,  0.7603, -2.0579,  1.7856, -3.9744,  0.1253,  2.0046,-0.6214, -2.1419]], device='cuda:0')
tensor([1], device='cuda:0')
--------------第3张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 0
tensor([[13.5426, -6.6119,  1.1040, -5.3236, -6.8938,  2.2850,  3.4611, -2.5504,-0.0393, -0.4295]], device='cuda:0')
tensor([0], device='cuda:0')
--------------第4张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 4
tensor([[-2.2200, -3.2653, -3.9221, -7.7044, 12.9432, -5.3840, -0.3826,  1.3231,-2.7672,  6.4946]], device='cuda:0')
tensor([4], device='cuda:0')
--------------第5张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 1
tensor([[-2.1054, 10.7492, -0.0951, -2.4075,  3.1492, -5.6566, -1.9704,  4.0755,-0.5151, -0.8886]], device='cuda:0')
tensor([1], device='cuda:0')
--------------第6张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 4
tensor([[-7.6890,  0.1027, -5.0862, -4.9066, 10.8449, -1.0489, -2.3638,  2.8111,4.3393,  3.6312]], device='cuda:0')
tensor([4], device='cuda:0')
--------------第7张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 9
tensor([[-7.3084, -1.4986, -1.5550,  2.7101,  2.4086,  2.5202, -7.5940,  0.3350,2.0277,  7.9211]], device='cuda:0')
tensor([9], device='cuda:0')
--------------第8张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 5
tensor([[ -0.5667, -15.0522,  -2.6193,  -2.5653,  -1.2110,  13.1138,   5.9642,-7.2953,   4.0217,   1.9577]], device='cuda:0')
tensor([5], device='cuda:0')
--------------第9张图片--------------
图片尺寸 torch.Size([1, 28, 28])
标签 9
tensor([[-2.5326, -8.8390, -2.7222, -2.3805,  2.1510, -1.2283, -7.7543,  5.7798,5.2312, 11.3871]], device='cuda:0')
tensor([9], device='cuda:0')Process finished with exit code 0
s

说明准确率都还挺高的,这么模型训练的还可以


好了,这就是minist数据集训练了

写完这个就算是入门了

b站上有一个pytorch很好的入门视频

土堆pytorch入门

讲的很不错的,学完就能入门了

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com