pytorch mnist 最简单示例代码

导入库

from typing import *

import os
import sys
import random
import pickle
import logging
from datetime import datetime

try:
    from tqdm.notebook import tqdm as _tqdm
except:
    _tqdm = lambda x : x

import numpy as np
from matplotlib import pyplot as plt

import torch as t
import torch.nn as nn
import torch.optim as optim

import torchvision as tv
from torch.utils.data import DataLoader

准备数据集

data_dir = '.'
batch_size = 256
train_dataset = tv.datasets.MNIST(root=data_dir, train=True, transform=tv.transforms.ToTensor(), download=True)
test_dataset = tv.datasets.MNIST(root=data_dir, train=False, transform=tv.transforms.ToTensor(), download=True)

logging 和 seed

logging.basicConfig(
    format='%(asctime)s %(levelname)s %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    level=logging.INFO,
    stream=sys.stdout,
    force=True,
)

_SEED = 0
random.seed(_SEED)
np.random.seed(_SEED)
t.manual_seed(_SEED)
t.cuda.manual_seed_all(_SEED)
# t.backends.cudnn.benchmark = False # burden for performance
# t.backends.cudnn.deterministic = True # burden for performance
# t.use_deterministic_algorithms(True) # deprecated in cuda 11+

assert random.random() == 0.662743203727557, 'different random generator'
assert np.random.rand() == 0.11717829958543136, 'different random generator'
assert t.rand((1,)).item() == 0.997646152973175, 'different random generator'

if t.cuda.is_available():
    for i in range(t.cuda.device_count()):
        logging.info(t.cuda.get_device_properties(i))
    logging.warning('using cuda:0')
    device = t.