layer 별 learning rate 할당 방법

 

import torch

import torch.nn as nn

import torchvision.models as models




class CNN(nn.Module):

    def __init__(self, ):

        super(CNN, self).__init__()

        self.backbone = models.resnet18(pretrained=True)

        for p in self.backbone.parameters():

            p.requires_grad = True

        self.fc = nn.Linear(512, 2)

    def forward(self, x):

        x = self.backbone.conv1(x)

        x = self.backbone.bn1(x)

        x = self.backbone.relu(x)

        x = self.backbone.maxpool(x)



        x = self.backbone.layer1(x)

        x = self.backbone.layer2(x)

        x = self.backbone.layer3(x)

        x = self.backbone.layer4(x)



        x = self.backbone.avgpool(x)

        x = torch.flatten(x, start_dim=1)

        x = self.fc(x)

        return x



net = CNN().to('cuda')

optimizer = torch.optim.Adam([{'params': net.fc.parameters(), 'lr':1e-2}, 

                              {'params': net.backbone.parameters(), 'lr': 1e-2}])



for i in range(10):

  for param_group in optimizer.param_groups:

      print(param_group['lr'])

  input = torch.rand((7, 3, 224, 224)).to('cuda')

  output = net(input)



  loss = torch.mean(output)

  optimizer.zero_grad()

  loss.backward()

  optimizer.step()

  

  print(loss)