import numpy as np
import torch
import torch. nn as nn
import torchvision
import torchvision. transforms as transforms
input_size = 28 * 28
num_classes = 10
num_epochs = 10
batch_size = 100
learning_rate = 0.001
train_dataset = torchvision. datasets. MNIST( root= '../../data' ,
train= True ,
transform= transforms. ToTensor( ) ,
download= True )
test_dataset = torchvision. datasets. MNIST( root= '../../data' ,
train= False ,
transform= transforms. ToTensor( ) )
train_loader = torch. utils. data. DataLoader( dataset= train_dataset,
batch_size= batch_size,
shuffle= True )
test_loader = torch. utils. data. DataLoader( dataset= test_dataset,
batch_size= batch_size,
shuffle= False )
model = nn. Linear( input_size, num_classes)
criterion = nn. CrossEntropyLoss( )
optimizer = torch. optim. SGD( model. parameters( ) , lr= learning_rate)
total_step = len ( train_loader)
for epoch in range ( num_epochs) :
for i, ( images, labels) in enumerate ( train_loader) :
images = images. reshape( - 1 , input_size)
outputs = model( images)
loss = criterion( outputs, labels)
optimizer. zero_grad( )
loss. backward( )
optimizer. step( )
if ( i+ 1 ) % 100 == 0 :
print ( 'Epoch [{}/{}], Step [{}/{}], Loss:{:.4f}' .
format ( epoch+ 1 , num_epochs, i+ 1 , total_step, loss. item( ) ) )
Epoch [1/10], Step [100/600], Loss:2.1877
Epoch [1/10], Step [200/600], Loss:2.0952
Epoch [1/10], Step [300/600], Loss:2.0192
Epoch [1/10], Step [400/600], Loss:1.8990
Epoch [1/10], Step [500/600], Loss:1.8696
Epoch [1/10], Step [600/600], Loss:1.7758
Epoch [2/10], Step [100/600], Loss:1.6857
Epoch [2/10], Step [200/600], Loss:1.6319
Epoch [2/10], Step [300/600], Loss:1.6608
Epoch [2/10], Step [400/600], Loss:1.5398
Epoch [2/10], Step [500/600], Loss:1.4866
Epoch [2/10], Step [600/600], Loss:1.5251
Epoch [3/10], Step [100/600], Loss:1.3724
Epoch [3/10], Step [200/600], Loss:1.3550
Epoch [3/10], Step [300/600], Loss:1.4124
Epoch [3/10], Step [400/600], Loss:1.3413
Epoch [3/10], Step [500/600], Loss:1.1988
Epoch [3/10], Step [600/600], Loss:1.2809
Epoch [4/10], Step [100/600], Loss:1.2419
Epoch [4/10], Step [200/600], Loss:1.1801
Epoch [4/10], Step [300/600], Loss:1.2190
Epoch [4/10], Step [400/600], Loss:1.1891
Epoch [4/10], Step [500/600], Loss:1.1165
Epoch [4/10], Step [600/600], Loss:1.1057
Epoch [5/10], Step [100/600], Loss:0.9912
Epoch [5/10], Step [200/600], Loss:1.1478
Epoch [5/10], Step [300/600], Loss:0.9327
Epoch [5/10], Step [400/600], Loss:0.9662
Epoch [5/10], Step [500/600], Loss:0.8517
Epoch [5/10], Step [600/600], Loss:0.9587
Epoch [6/10], Step [100/600], Loss:0.9835
Epoch [6/10], Step [200/600], Loss:0.9946
Epoch [6/10], Step [300/600], Loss:0.8951
Epoch [6/10], Step [400/600], Loss:0.9013
Epoch [6/10], Step [500/600], Loss:0.9931
Epoch [6/10], Step [600/600], Loss:0.8686
Epoch [7/10], Step [100/600], Loss:0.9099
Epoch [7/10], Step [200/600], Loss:0.8394
Epoch [7/10], Step [300/600], Loss:0.9348
Epoch [7/10], Step [400/600], Loss:0.7706
Epoch [7/10], Step [500/600], Loss:1.0090
Epoch [7/10], Step [600/600], Loss:0.8198
Epoch [8/10], Step [100/600], Loss:0.8123
Epoch [8/10], Step [200/600], Loss:0.8300
Epoch [8/10], Step [300/600], Loss:0.8192
Epoch [8/10], Step [400/600], Loss:0.7725
Epoch [8/10], Step [500/600], Loss:0.8281
Epoch [8/10], Step [600/600], Loss:0.8399
Epoch [9/10], Step [100/600], Loss:0.8856
Epoch [9/10], Step [200/600], Loss:0.7889
Epoch [9/10], Step [300/600], Loss:0.7205
Epoch [9/10], Step [400/600], Loss:0.7489
Epoch [9/10], Step [500/600], Loss:0.8840
Epoch [9/10], Step [600/600], Loss:0.6865
Epoch [10/10], Step [100/600], Loss:0.7590
Epoch [10/10], Step [200/600], Loss:0.7337
Epoch [10/10], Step [300/600], Loss:0.7751
Epoch [10/10], Step [400/600], Loss:0.6656
Epoch [10/10], Step [500/600], Loss:0.6606
Epoch [10/10], Step [600/600], Loss:0.7351
with torch. no_grad( ) :
correct = 0
total = 0
for images, labels in test_loader:
images = images. reshape( - 1 , input_size)
outputs = model( images)
_, predicted = torch. max ( outputs. data, 1 )
total += labels. size( 0 )
correct += ( predicted== labels) . sum ( )
print ( 'Accuracy of the model on the 10000 test images:{} %' .
format ( 100 * correct/ total) )
Accuracy of the model on the 10000 test images:85.55999755859375 %
torch. save( model. state_dict( ) , 'model.ckpt' )