利用pytorch随便搭建一个卷积神经网络,实现mnist的分类任务,准确率达到100%。
友情提示:pytorch推荐版本1.0及以上,低版本的Tensor变量无法向Variable变量一样进行backward操作,若未安装GPU版本的pytorch,记得删掉代码中所有的 .cuda() 。
( .cuda() 为网络模型、输入图片、输入标签移至GPU的操作)
注:若pytorch数据集下载速度过慢,可以参考此博客解决https://blog.youkuaiyun.com/weixin_44414948/article/details/109756003
代码如下(训练时可将数据可视化模块代码注释掉):
import torch
import torchvision
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.utils.data as Data
EPOCH=1#训练整批数据多少次
BATCH_SIZE=64#每次批数的数据量
LR=0.001#学习率,学习率的设置直接影响着神经网络的训练效果
train_data=torchvision.datasets.MNIST(#训练数据
root= './mnist_data/',
train=True,
transform=torchvision.transforms.ToTensor(),
download=True
)
test_data = torchvision.datasets.MNIST(
root='./mnist_data/',
train=False,
transform=torchvision.transforms.ToTensor(),
download=True
)
# 批量加载
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = Data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE,