Pytorch的Helloworld

最新推荐文章于 2025-05-24 15:12:56 发布

轶张

最新推荐文章于 2025-05-24 15:12:56 发布

阅读量1.2k

点赞数

CC 4.0 BY-SA版权

分类专栏： Lab

本文链接：https://blog.youkuaiyun.com/qq_35941161/article/details/103637945

可直接运行。
受益良多

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

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

train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=BATCH_SIZE, shuffle=True)

test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data', train=False, transform=transforms