基于Pytorch中MNIST数据集实现手写数字识别
环境要求:python且安装torch、torchvision库
数据集准备
导入需要使用模块
torchvision中有本次使用的手写数字照片数据集
transforms将数据集照片进行格式化转换。如:照片转化为tensor格式、进行归一化处理
nn为torch模型板块、optim为模型优化板块、utils为torch工具板块
import torch
import torchvision
import torch. utils as utils
import torch. optim as optim
import torch. nn as nn
import torchvision. transforms as transforms
对数据进行处理以及分批次
掉用torchvision.datasets.MNIST获取数据集
将数据集按照transform进行格式化处理
将数据集按照BATCH_SIZE大小进行分批次(shuffle=True,打乱随机)
BATCH_SIZE = 10
transform = transforms. Compose( [
transforms. ToTensor( ) ,
transforms. Normalize( mean= 0.1307 , std= 0.3081 )
] )
device = ( "cuda:0" if torch. cuda. is_available( ) else "cpu" )
def dataload ( train= True ) :
MNIST_data = torchvision. datasets. MNIST( "./DATA/" , train = train , download= True , transform= transform)
train_dataloead = utils. data. DataLoader( dataset= MNIST_data, batch_size= BATCH_SIZE, shuffle= True )
return train_dataloead
train_dataload = dataload( )
构建模型
继承nn.Module构建自己的识别模型
搭建一个3层隐藏层,以及为前两层加了非线性激活函数
由于数据集为2828图像,所以将数据集压缩成一个x.size(0),28 28二维张量
class Model_MNIST ( nn. Module) :
def __init__ ( self) :
super ( Model_MNIST, self) . __init__( )
self. fx1 = nn. Linear( 28 * 28 , 392 )
self. fx2 = nn. Linear( 392 , 98 )
self. fx3 = nn. Linear( 98 , 10 )
self. relu = nn. ReLU( )
def