3.1 利用MNIST数据集做手写数字识别

最新推荐文章于 2024-05-21 14:45:00 发布

羊藤枝

最新推荐文章于 2024-05-21 14:45:00 发布

阅读量321

点赞数 1

分类专栏： pytorch 机器学习深度学习文章标签：深度学习

本文链接：https://blog.youkuaiyun.com/qq_30614451/article/details/93913440

版权

这篇博客介绍了如何利用MNIST数据集进行手写数字识别，通过设定超参数，导入PyTorch的MNIST数据包，并定义了一个卷积神经网络模型。模型实例化后，将其转移到GPU进行计算。文章还详细阐述了训练过程，将所有操作封装在一个训练函数中。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

3.1 利用MNIST数据集做手写数字识别

from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

首先，设置超参数：

BATCH_SIZE=512 #大概需要2G的显存
EPOCHS=20 # 总共训练批次
# 让torch判断是否使用GPU，建议使用GPU环境，因为会快很多
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

从pytorch里面导入此次使用的mnist数据包：

transform = transforms.Compose(
    [transforms.ToTensor(),
      transforms.Normalize((0.1307,),(0.3081,))])



trainset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, 
                                         download=True)
train_loader = torch.utils.data.DataLoader(trainset,batch_size=BATCH_SIZE,shuffle=True)


testset = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, 
                                         download=True)
test_loader = torch.utils.<