神经网络模型对手写数字的识别
import torch
from torch import nn
from torch. utils. data import DataLoader
from torchvision import datasets
from torchvision. transforms import ToTensor
"""
MNIST包含70,000张手写数字图像:60,000张用于训练,10,000张用于测试。
图像是灰度的,28x28像素的,并且居中的,以减少预处理和加快运行。
"""
""" 下载训练数据集 (包含训练数据+标签)"""
training_data = datasets. MNIST(
root= 'data' ,
train= True ,
download= True ,
transform= ToTensor( )
)
""" 下载测试数据集(包含训练图片+标签)"""
test_data = datasets. MNIST(
root= 'data' ,
train= False ,
download= True ,
transform= ToTensor( )
)
print ( len ( training_data) )
""" 展示手写字图片 """
from matplotlib import pyplot as plt
figure = plt. figure( )
for i in range ( 9 ) :
img, label = training_data[ i + 59000 ]
figure. add_subplot( 3 , 3 , i + 1 )
plt. title( label)
plt. axis( "off" )
plt. imshow( img. squeeze( ) , cmap= "gray" )
a = img. squeeze( )
plt. show( )
training_dataloader = DataLoader( training_data, batch_size= 64 )
test_dataloader = DataLoader( test_data, batch_size= 64 )
for X, y in test_dataloader:
print ( f"Shape of X [N, C, H, W]: { X. shape} " )
print ( f"Shape of y: { y. shape} { y. dtype} " )
break
""" 判断当前设备是否支持GPU,其中mps是苹果m系列芯片的GPU """
device = "cuda" if torch. cuda. is_available( ) else "mps" if torch. backends. mps. is_available( ) else "cpu"
print ( f"Using { device} device" )
class NeuralNetwork ( nn. Module) :
def __init__ ( self) :
super ( ) . __init__( )
self. flatten = nn. Flatten( )
self. hidden1 = nn. Linear( 28 * 28 , 256 )
self. hidden2 = nn. Linear( 256 , 128 )
self. hidden3 = nn. Linear( 128 , 256 )
self. hidden4 = nn. Linear( 256 , 128 )
self. out = nn. Linear( 128 , 10 )
def forward ( self, x) :
x = self. flatten( x)
x = self. hidden1( x)
x = torch. sigmoid( x)
x = self. hidden2( x)
x = torch. sigmoid( x)
x = self. hidden3( x)
x = torch. sigmoid( x)
x = self. hidden4( x)
x = torch. sigmoid( x)
x = self. out( x)
return x
model = NeuralNetwork( ) . to( device)
print ( model)
def train ( dataloader, model, loss_fn, optimizer) :
model. train( )
batch_size_num = 1
for X, y in dataloader:
X, y = X. to( device) , y. to( device)
pred = model. forward( X)
loss = loss_fn( pred, y)
optimizer. zero_grad( )
loss. backward( )
optimizer. step( )
loss_value = loss. item( )
if batch_size_num % 200 == 0 :
print ( f"loss: { loss_value: >7f } [number: { batch_size_num} ]" )
batch_size_num += 1
def test ( dataloader, model, loss_fn) :
size = len ( dataloader. dataset)
num_batches = len ( dataloader)
model. eval ( )
test_loss, correct = 0 , 0
with torch. no_grad( ) :
for X, y in dataloader:
X, y = X. to( device) , y. to( device)
pred = model. forward( X)
test_loss += loss_fn( pred, y) . item( )
correct += ( pred. argmax( 1 ) == y) . type ( torch. float ) . sum ( ) . item( )
a = ( pred. argmax( 1 ) == y)
b = ( pred. argmax( 1 ) == y) . type ( torch. float )
test_loss /= num_batches
correct /= size
print ( f"Test result: \n Accuracy: { ( 100 * correct) } %, Avg loss: { test_loss} " )
loss_fn = nn. CrossEntropyLoss( )
optimizer = torch. optim. Adam( model. parameters( ) , lr= 0.01 )
epochs = 10
for e in range ( epochs) :
print ( f"Epoch { e + 1 } \n" )
train( training_dataloader, model, loss_fn, optimizer)
print ( "Done!" )
test( test_dataloader, model, loss_fn)
展示的手写数字图片如下: 模型结构如下: 训练结果如下: 共有10轮训练 测试结果如下: