How make two model‘s block with same structure share weight？

Question

There two model for example class model1(nn.Module) and class model2(nn.Module). Encoder block is in the two model with same structure. So how make two model's encoder block share weight?

Make two model's encoder block share weight.

For example:

class Model1(nn.Module):
    def __init__(self, n_channels=3):
        super().__init__()

        # shared weight in encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(n_channels, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )


    def forward(self, x):
        x = self.encoder(x)
        return 

class Model2(nn.Module):
    def __init__(self, n_channels=3):
        super().__init__()

        # shared weight in encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(n_channels, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )


    def forward(self, x):
        x = self.encoder(x)
        return 

How can I make the encoder block in Model1 and Model2 share weight? (The encoder block must be in different modules.)

Answer 1

The implementation of encoder layers in model1 and model2 should be same.

here's a simple model

class Model(nn.Module):
def __init__(self, n_channels=3):
    super().__init__()

    # shared weight in encoder
    self.encoder = nn.Sequential(
        nn.Conv2d(n_channels, 64, kernel_size=3, padding=1),
        nn.BatchNorm2d(64),
        nn.ReLU(inplace=True),
    )

    self.head1 = nn.Conv2d(64, 10, kernel_size=3, padding=1) # task 1
    self.head2 = nn.Conv2d(64, 20, kernel_size=3, padding=1) # task 2
    

def forward(self, x):
    x = self.encoder(x)

    out1 = self.head1(x)    # pass through same input
    out2 = self.head2(x)    # pass through same input

    return