How to copy specific layer strucutres and weights to a new model

Hi, how can I copy specific layer structures and weights to a new model? I need to use Resnet but with additional layers and operations. What would be the best practice?

here’s my sample code

class Encoder(nn.Module):
    def __init__(self, args1, args2):
        super().__init__()

        resnet = models.resnet50(pretrained=True)
        # conv1, bn, relu, maxpool
        self.stem = resnet.stem
        # resnet block 1,2
        self.layer1 = resnet.layer1
        self.layer2 = resnet.layer2
        # some new layers
        self.layer3 =  MultiInputBlock(args)
    def forward(self, x, y):
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        if some_conditions:
           x = self.layer3(x, y)
        return x

Your code looks alright and should work assuming you fix the attribute names.
Alternatively, you could also create a custom model by deriving from a resnet implementation in case that’s easier or cleaner.

1 Like