Modifying stem and use pretrained weights

1

If I modify the stem() for torchvision models, will I be able to use the pretrained wieghts?

I am changing the input layer channels:

class modifybasicstem(nn.Sequential):
    """The default conv-batchnorm-relu stem
    """
    def __init__(self):
        super(modifybasicstem, self).__init__(
            nn.Conv3d(1, 64, kernel_size=(3, 3, 3), stride=(1, 2, 2),
                      padding=(1, 3, 3), bias=False),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True))
#
model = torchvision.models.video.mc3_18(pretrained=False)
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, num_classes)
)

model = torchvision.models.video.r3d_18(pretrained=True)
model.stem = modifybasicstem()
model.to('cuda:0')

Thankyou.

2

Yes, your code would replace the .stem module only while all other layers would still use their pretrained parameters.

1 Like