pytorch使用GPU训练MNIST数据集

「已注销」

于 2019-02-24 12:44:13 发布

阅读量4.3k

点赞数 2

分类专栏：机器学习

本文链接：https://blog.youkuaiyun.com/nerdcfl/article/details/87902011

版权

本文记录了使用PyTorch在GPU上训练MNIST数据集的过程，参考了莫凡博客的内容，主要包括数据加载、模型构建、损失计算及反向传播等步骤。

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

参考莫凡博客进行MNIST数据集的训练，临时记录所使用的代码。


import torch
import torch.nn as nn
import torch.utils.data as Data
import  torchvision
import matplotlib.pyplot as plt

torch.manual_seed(1)

EPOCH = 1
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = True

train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST,
)

test_data = torchvision.datasets.MNIST(root='./mnist', train=False)

# 批处理
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=