Modifying stem and use pretrained weights

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),
model =
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, num_classes)

model =
model.stem = modifybasicstem()'cuda:0')


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

1 Like