Pytorch_NNdemo002

Pytorch_NNdemo002.py

# 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):
    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()