MASKRCNN, Error while loading the model

in get_instance_segmentation_model(num_classes)
24 model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True, progress=True,
25 num_classes=num_classes, pretrained_backbone=True,
—> 26 trainable_backbone_layers=3)
27
28

/usr/local/lib/python3.6/site-packages/torchvision/models/detection/mask_rcnn.py in maskrcnn_resnet50_fpn(pretrained, progress, num_classes, pretrained_backbone, **kwargs)
310 pretrained_backbone = False
311 backbone = resnet_fpn_backbone(‘resnet50’, pretrained_backbone)
–> 312 model = MaskRCNN(backbone, num_classes, **kwargs)
313 if pretrained:
314 state_dict = load_state_dict_from_url(model_urls[‘maskrcnn_resnet50_fpn_coco’],

TypeError: init() got an unexpected keyword argument ‘trainable_backbone_layers’

The trainable_backbone_layers argument was added ~2months ago in this PR, so you might need to update torchvision.

Hello, I just upgraded to torch 1.5.1 and torchvision 0.6.1 but I am still getting this error. my code is bellow

import torch
import torchvision
import torchvision.models.detection.mask_rcnn
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

print(torch.__version__)
print(torchvision.__version__)

1.5.1
0.6.1

def get_model(num_classes, pretrained=True):
    
    # load an instance segmentation model pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=pretrained, trainable_backbone_layers=5)

    # get the number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # now get the number of input features for the mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    return model

model = get_model(2)

TypeError Traceback (most recent call last)
in
----> 1 model = get_model(2)

in get_model(num_classes, pretrained)
2
3 # load an instance segmentation model pre-trained on COCO
----> 4 model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=pretrained, trainable_backbone_layers=5)
5
6 # get the number of input features for the classifier

/opt/conda/lib/python3.6/site-packages/torchvision/models/detection/mask_rcnn.py in maskrcnn_resnet50_fpn(pretrained, progress, num_classes, pretrained_backbone, **kwargs)
316 pretrained_backbone = False
317 backbone = resnet_fpn_backbone(‘resnet50’, pretrained_backbone)
–> 318 model = MaskRCNN(backbone, num_classes, **kwargs)
319 if pretrained:
320 state_dict = load_state_dict_from_url(model_urls[‘maskrcnn_resnet50_fpn_coco’],

TypeError: init() got an unexpected keyword argument ‘trainable_backbone_layers’