logo

How to use optimizers in PyTorch

gcptutorials.com 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.Tanh(),
    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())

    sgd_optimizer.zero_grad()

    loss.backward()

    sgd_optimizer.step()

    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 

    

sgd_optimizer.step()

    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

    • Category: PyTorch