pytorch学习笔记-实践-CNN-minst手写数字识别

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

'1.2.0'

首先定义一些超参数

BATCH_SIZE = 512
EPOCHS = 10
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 让torch判断是否使用GPU，建议使用GPU环境，因为会快很多
print(DEVICE)

cpu

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

<torch.utils.data.dataloader.DataLoader object at 0x00000238F3EC1F88>

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