Search code examples
pythonpytorch

Can i initialize optimizer before changing the last layer of my model

Say I want to change the last layer of my code but my optimizer is defined on the top of my scripts, what is a better practise?

batch_size = 8
learning_rate = 2e-4
num_epochs = 100

cnn = models.resnet18(weights='DEFAULT')
loss_func = nn.CrossEntropyLoss()
optimizer = optim.Adam(cnn.parameters(), lr=learning_rate)

def main(): 
  dataset= datasets.ImageFolder(root=mydir)
  loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
  num_classes = len(dataset.classes)
  num_ftrs = cnn.fc.in_features
  cnn.fc = nn.Linear(num_ftrs, num_classes)
  
  # Do i need to reinitialize the optimizer again here? 

if __name__ == '__main__':
  main()
Solution

  • No, the optimizer would not be aware of the weights for the new last layer unless you initialize or reinitilize the optimizer after modifying the model.

    You should define the optimizer after creating the model and modifying any layers as needed.