Help me load a pretrained model

1

I have this model that I trained.

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=28, kernel_size=(6,6))
        self.conv2 = nn.Conv2d(in_channels=28, out_channels=56, kernel_size=(3,3))
        
        return x

Now I build another model

class CNN2(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=28, kernel_size=(6,6))
        self.fc2 = .....
        
        return x

How can I load the trained weights of model 1 (conv1) into model 2 (conv1) ? I do not need (conv2). And then freeze the weights of conv1

2

You could try the following code snippet:

model1 = CNN()
model2 = CNN2()

# Copy weights
model2.conv1.weight.data.copy_(model1.conv1.weight.data)
print(torch.abs(model2.conv1.weight - model1.conv1.weight).sum())
> tensor(0.)

# Freeze conv1
model2.conv1.weight.requires_grad_(False)
print(model2.conv1.weight.requires_grad)
> False
1 Like