PyTorch如何進行Transfer Learning

一般來說，進行Transfer Learning時，都會使用PyTorch官方在ImageNet上的Pre-trained model，並且將模型的參數進行Freeze，再搭建自定義的全連接層(FC Layer)進行模型訓練。

2022/03/16更新：PyTorch使用Pre-trained model進行Transfer Learning

例如官方提供的這個範例：

model_conv = torchvision.models.resnet18(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False

# Parameters of newly constructed modules have requires_grad=True by default
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()

# Observe that only parameters of final layer are being optimized as
# opposed to before.
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)

官方提供的範例相當的簡單，但也相當的侷限。

如果我們想要再對模型的forward進行修改或是新增神經網路層(Layer)，就必須參考這篇討論串中的做法：

class MyCustomResnet18(nn.Module):
    def __init__(self, pretrained=True):
        super().__init__()
        
        #特別注意的是，這個resnet18的變數，並不一定要是官方預訓練的模型
        #也可以是自定義的神經網路架構，增加了使用上的彈性
        resnet18 = models.resnet18(pretrained=pretrained)
        self.features = nn.ModuleList(resnet18.children())[:-1]
        self.features = nn.Sequential(*self.features)
        in_features = resnet18.fc.in_features
        self.fc0 = nn.Linear(in_features, 256)
        self.fc0_bn = nn.BatchNorm1d(256, eps = 1e-2)
        self.fc1 = nn.Linear(256, 256)
        self.fc1_bn = nn.BatchNorm1d(256, eps = 1e-2)
        
        for m in self.modules():
            if isinstance(m, nn.Linear):
                torch.nn.init.xavier_normal_(m.weight, gain = 1)
    def forward(self, input_imgs):
        #如果使用自定義神經網路架構時，要注意nn.Sequential搭配
        #nn.ModuleList後，可能會因為定義之架構屬多段組合而成之網路
        #多了一層額外的Sequential架構造成圖片送入網路時dim錯誤，因此可以改用:
        #out = self.features[0](x)
        #out = self.features[1](x)
        #的方式來解決。
        output = self.features(input_imgs)
        output = output.view(input_imgs.size(0), -1)
        output = self.fc0_bn(F.relu(self.fc0(output)))
        output = self.fc1_bn(F.relu(self.fc1(output)))
        return output

參考資料

Transfer Learning for Computer Vision Tutorial - PyTorch Tutorials 1.10.0+cu102 documentation

How can I replace the forward method of a predefined torchvision model with my customized forward…