学习Pytorch+Python之MNIST手写字体识别

 
 好买网  https://www.goodmai.com/ 

import matplotlib.pyplot as plt

import torch

import torch.nn as nn

import numpy as np

import torchvision.utils

from torchvision import datasets, transforms

from torch.autograd import Variable

import torch.utils.data

#判断是否能用GPU，如果能就用GPU，不能就用CPU

use_gpu = torch.cuda.is_available()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#数据转换，Pytorch的底层是tensor(张量)，所有用来训练的图像均需要转换成tensor

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

#下载数据集

data_train = datasets.MNIST(root="./data/", transform=transform, train=True, download=True)

data_test  = datasets.MNIST(root="./data/", transform=transform, train=False)

#加载数据集，批次大小为64，shuffle表示乱序

data_loader_train = torch.utils.data.DataLoader(dataset=data_train, batch_size=64, shuffle=True)

data_loader_test = torch.utils.data.DataLoader(dataset=data_test, batch_size=64, shuffle=True)

#创建模型即网络架构

class Model(nn.Module):