Code snippet for: https://wp.me/p5EUYy-fq

Pytorch NN Demo.py

n_input = 784
n_hidden = 200
n_output = 10

# Model
model = nn.Sequential(nn.Linear(n_input, 128),
                     nn.ReLU(),
                     nn.Linear(128, 64),
                     nn.ReLU(),
                     nn.Linear(64,10),
                     nn.LogSoftmax(dim=1))
# Defining gradient descent
learning_rate = 0.01
criterion = nn.NLLLoss()
optimiser = optim.SGD(model.parameters(), lr=learning_rate)

# Training the network
epochs = 5

for e in range(epochs):
    train_loss = 0
    for images, labels in train_loader:
        # Flattening an input image to vector
        images = images.view(images.shape[0], -1)
        # Initialising model parameters' gradient to zero
        optimiser.zero_grad()
        
        # Propogating forward: input -> hidden -> output
        output = model.forward(images)        
        # Calculating error at output
        loss = criterion(output, labels)
        # Propogating an error backward: output -> hidden -> input
        loss.backward()
        # Updating weight
        optimiser.step()
        
        # Recoring loss of error
        train_loss += loss.item()
        
    print(f'Training loss:{train_loss/len(train_loader)}')