What parameters do I change to train a pytorch model from scratch?

user_1185242 0
Viewed 1653

I followed this tutorial to train a pytorch model for instance segmentation: https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html

I would not like to train a model on entirely different data and classes, totally unrelated to COCO. What changes do I need to make to retrain the model. From my reading I'm guessing besides have the correct number of classes I just need to train this line:

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

to

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False)

But I notice there is another parameters: pretrained_backbone=True, trainable_backbone_layers=None should they be changed too?

pytorch deep-learning python
2 Answers

The function signature is

torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True, trainable_backbone_layers=3, **kwargs)

Setting pretrained=False will tell PyTorch not to download model pre-trained on COCO train2017. You want it as you're interested in training.

Usually, this is enough if you want to train on a different dataset.

When you set pretrained=False, PyTorch will download pretrained ResNet50 on ImageNet. And by default, it'll freeze first two blocks named conv1 and layer1. This is how it was done in Faster R-CNN paper which frooze the initial layers of pretrained backbone.

(Just print model to check its structure).

layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers]

Now, if you don't even want the first two layers to freeze, you can set trainable_backbone_layers=5 (done automatically, when you set pretrained_backbone=False), which will train the entire resnet backbone from scratch.

Check PR#2160.

user_62108070

From maskrcnn_resnet50_fpn document:

  • pretrained (bool) – If True, returns a model pre-trained on COCO train2017
  • pretrained_backbone (bool) – If True, returns a model with backbone pre-trained on Imagenet
  • trainable_backbone_layers (int) – number of trainable (not frozen) resnet layers starting from final block. Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.

So for training from scratch using:

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False, trainable_backbone_layers=5, num_classes=your_num_classes)

or:

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False, num_classes=your_num_classes)

because in source code of maskrcnn_resnet50_fpn:

if not (pretrained or pretrained_backbone):
    trainable_backbone_layers = 5
user_135095400
Related