基于pytorch的手写数字识别，算法优化（动量Momentum，梯度下降SGD）GPU加速

导入相关模块

import numpy as np
import torch
# 导入 pytorch 内置的 mnist 数据
from torchvision.datasets import mnist 
# 导入预处理模块
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# 导入nn及优化器
from torch import nn
import torch.optim as optim
import torch.nn.functional as F
​
# 数据可视化
import matplotlib.pyplot as plt
%matplotlib inline

定义参数

# 每次训练数据数量
train_batch_size = 64
# 每次测试数据数量
test_batch_size = 128
# 学习率
learning_rate = 0.01
# 训练轮次
num_epoches = 8
# 动量参数
momentum = 0.9

数据预处理

# 定义预处理函数，这些预处理依次放在Compose函数中  
# Normalize([0.5],[0.5])对张量进行归一化，这里两个0.5分别表示对张量进行归一化的全局平均值和方差。
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5],[0.5])])
# 下载数据，并对数据进行预处理  download参数控制是否需要下载，如果./data目录下已有MNIST，可选择False；
train_dataset = mnist.MNIST('./data',train=True,transform=transform,download=True)
test_dataset = mnist.MNIST('./data',train=False,transform=transform,download=True)
# dataloader是一个可迭代对象，可以和迭代器一样使用  用DataLoader得到生成器，这可节省内存
train_loader = DataLoader(train_dataset,batch_size=train_batch_size,shuffle=True)
test_loader = DataLoader(test_dataset,batch_size=test_batch_size,shuffle=False)

构建模型

class Net(nn.Module):
    """使用sequential构建网络，Sequential()函数的功能是将网络的层组合到一起"""
    def __in