How to add my custom layer? after load pre traind weight?

1

Does anyone know how to add my custom layer to a specific location after loading pre traind weights? without breaking the pre traind weight.

2

You could add the extra layer, manipulate the forward method as you wish, and load the state_dict passing strict=False.
Here is a small example:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(1, 1)
        self.fc2 = nn.Linear(1, 1)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

class MyModelModified(nn.Module):
    def __init__(self):
        super(MyModelModified, self).__init__()
        self.fc1 = nn.Linear(1, 1)
        self.fc2 = nn.Linear(1, 1)
        self.extra = nn.Linear(1, 1)

    def forwars(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.extra(x)
        return x

model = MyModel()
state_dict = model.state_dict()

model_mod= MyModelModified()
model_mod.load_state_dict(state_dict, strict=False)
> _IncompatibleKeys(missing_keys=['extra.weight', 'extra.bias'], unexpected_keys=[])

# Check for equal params
print(model.fc1.weight == model_mod.fc1.weight)
> True
print(model.fc1.bias == model_mod.fc1.bias)
> True
print(model.fc2.weight == model_mod.fc2.weight)
> True
print(model.fc2.bias == model_mod.fc2.bias)
> True