import torch
import torch. nn as nn
import torchvision. datasets as dsets
import torchvision. transforms as transforms
from torch. autograd import Variable
torch. manual_seed( 1 )
input_size = 784
hidden_size = 500
num_classes = 10
num_epoch = 5
batch_size = 100
learning_rate = 0.001
train_dataset = dsets. MNIST( root= './data' ,
train = True ,
transform = transforms. ToTensor( ) ,
download = True )
test_dataset = dsets. 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 = 1 , shuffle = False )
class Net ( nn. Module) :
def __init__ ( self, input_size, hidden_size, num_classes) :
super ( Net, self) . __init__( )
self. fc1 = nn. Linear( input_size, hidden_size)
self. relu = nn. ReLU( )
self. fc2 = nn. Linear( hidden_size, num_classes)
def forward ( self, x) :
out = self. fc1( x)
out = self. relu( out)
out = self. fc2( out)
return out
net = Net( input_size, hidden_size= hidden_size, num_classes= num_classes)
print ( net)
Net(
(fc1): Linear(in_features=784, out_features=500, bias=True)
(relu): ReLU()
(fc2): Linear(in_features=500, out_features=10, bias=True)
)
criterion = nn. CrossEntropyLoss( )
optimizer = torch. optim. Adam( net. parameters( ) , lr= learning_rate)
for epoch in range ( num_epoch) :
for i, ( images, labels) in enumerate ( train_loader) :
images = Variable( images. view( - 1 , 28 * 28 ) )
labels = Variable( labels)
optimizer. zero_grad( )
outputs = net( images)
loss = criterion( outputs, labels)
loss. backward( )
optimizer. step( )
if ( i+ 1 ) % 100 == 0 :
print ( 'epoch:' , epoch, 'step:' , i, 'loss:' , loss. item( ) )
epoch: 0 step: 99 loss: 0.3248719274997711
epoch: 0 step: 199 loss: 0.3053732216358185
epoch: 0 step: 299 loss: 0.2878018319606781
epoch: 0 step: 399 loss: 0.1890110969543457
epoch: 0 step: 499 loss: 0.23649799823760986
epoch: 0 step: 599 loss: 0.15793423354625702
epoch: 1 step: 99 loss: 0.09613646566867828
epoch: 1 step: 199 loss: 0.09829027950763702
epoch: 1 step: 299 loss: 0.14937524497509003
epoch: 1 step: 399 loss: 0.0760403499007225
epoch: 1 step: 499 loss: 0.08814708888530731
epoch: 1 step: 599 loss: 0.04184893146157265
epoch: 2 step: 99 loss: 0.06356063485145569
epoch: 2 step: 199 loss: 0.051955245435237885
epoch: 2 step: 299 loss: 0.1010856181383133
epoch: 2 step: 399 loss: 0.06651352345943451
epoch: 2 step: 499 loss: 0.16046684980392456
epoch: 2 step: 599 loss: 0.09681060165166855
epoch: 3 step: 99 loss: 0.04019105061888695
epoch: 3 step: 199 loss: 0.08993328362703323
epoch: 3 step: 299 loss: 0.06461403518915176
epoch: 3 step: 399 loss: 0.10949535667896271
epoch: 3 step: 499 loss: 0.042061515152454376
epoch: 3 step: 599 loss: 0.1040988340973854
epoch: 4 step: 99 loss: 0.013200558722019196
epoch: 4 step: 199 loss: 0.09780997037887573
epoch: 4 step: 299 loss: 0.0756884515285492
epoch: 4 step: 399 loss: 0.006697341799736023
epoch: 4 step: 499 loss: 0.05417792499065399
epoch: 4 step: 599 loss: 0.037775710225105286
correct = 0
total = 0
for images, labels in test_loader:
images = Variable( images. view( - 1 , 28 * 28 ) )
outputs = net( images)
_, predicted = torch. max ( outputs. data, 1 )
total += labels. size( 0 )
correct += ( predicted == labels) . sum ( )
print ( total)
print ( ( correct/ total) . item( ) )
10000
0.9787999987602234