Use Face Alignment Detector as backbone for pytorch faster rcnn

[toti-bz]

I am trying to use face alignment network as backbone for faster rcnn pytorch pretrained model. I am using the net_s3fd structure as a feature extraction network and this is how I implement it. `def create_model(num_classes):

fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device='cuda', face_detector='sfd')

s3fd_detector = fa.face_detector.face_detector

conv1_1 = s3fd_detector.conv1_1
conv1_2 = s3fd_detector.conv1_2
conv2_1 = s3fd_detector.conv2_1
conv2_2 = s3fd_detector.conv2_2
conv3_1 = s3fd_detector.conv3_1
conv3_2 = s3fd_detector.conv3_2
conv3_3 = s3fd_detector.conv3_3
conv4_1 = s3fd_detector.conv4_1
conv4_2 = s3fd_detector.conv4_2
conv4_3 = s3fd_detector.conv4_3
conv5_1 = s3fd_detector.conv5_1
conv5_2 = s3fd_detector.conv5_2
conv5_3 = s3fd_detector.conv5_3
fc6 = s3fd_detector.fc6
fc7 = s3fd_detector.fc7
conv6_1 = s3fd_detector.conv6_1
conv6_2 = s3fd_detector.conv6_2
conv7_1 = s3fd_detector.conv7_1
conv7_2 = s3fd_detector.conv7_2

s3fd = nn.Sequential(OrderedDict([
    ('conv1_1', conv1_1), ('conv1_2', conv1_2), 
    ('conv2_1', conv2_1), ('conv2_2', conv2_2),
    ('conv3_1', conv3_1), ('conv3_2', conv3_2), ('conv3_3', conv3_3),
    ('conv4_1', conv4_1), ('conv4_2', conv4_2), ('conv4_3', conv4_3),
    ('conv5_1', conv5_1), ('conv5_2', conv5_2), ('conv5_3', conv5_3), 
    ('fc6', fc6), ('fc7', fc7),
    ('conv6_1', conv6_1), ('conv6_2', conv6_2),
    ('conv7_1', conv7_1), ('conv7_2', conv7_2) 
]))

backbone = s3fd
backbone.out_channels = 256

anchor_generator = AnchorGenerator(
    sizes=((32, 64, 128, 256, 512),),
    aspect_ratios=((0.5, 1.0, 2.0),)
)

roi_pooler = torchvision.ops.MultiScaleRoIAlign(
    featmap_names=['0'],
    output_size=7,
    sampling_ratio=2
)

# Final Faster RCNN model.
model = FasterRCNN(
    backbone=backbone,
    num_classes=num_classes,
    rpn_anchor_generator=anchor_generator,
    box_roi_pool=roi_pooler
)
print(model)
return model`

The problem is when I start training on google colab or kaggle I get "Cuda out of memory" error. Is there a way to use this network as the feature extraction network with the pretrained weights other than the one I'm using?

Thanks