How to use optimizers in PyTorch

PyTorch provides torch.optim package for implementing optimization algorith for a neural network. torch.optim supports commonly used optimizers, that can be directly invoked as torch.optim.

Steps for using a optimizer

  • Construct the optimizer by providing parameters
  • Update the parameters with step() method
  • Let's see how to use optimizer with the help of below code snippet.

    import torch
    batch_size, input_dim, hidden_dim, out_dim = 32, 100, 100, 10
    input_tensor = torch.randn(batch_size, input_dim)
    output_tensor = torch.randn(batch_size, out_dim)
    model = torch.nn.Sequential(
        torch.nn.Linear(input_dim, hidden_dim),
        torch.nn.Linear(hidden_dim, out_dim),
    loss_function = torch.nn.MSELoss(reduction='sum')
    lr = 1e-5
    sgd_optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    for i in range(200):
        predicted_value = model(input_tensor)
        loss = loss_function(predicted_value, output_tensor)
        print(i, loss.item())

    Above code snippet uses SGD optimizer for training the network.

    Constructing sgd optimizer

    sgd_optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    Updating parameters of sgd optimizer


    Similarly other optimizers can be used with torch.optim

  • torch.optim.Adadelta: Implements Adadelta algorithm
  • torch.optim.Adagrad: Implements Adagrad algorithm
  • torch.optim.Adam: Implements Adam algorithm.
  • torch.optim.Adamax: Implements Adamax algorithm
  • Categories: PyTorch