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.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
  • Categories: PyTorch