Reuse existing model

1

I am trying to reuse a model (transfer learning). I want to remove the last layer (and will later add another layer)

My existing model is

 class Model_TL_A(nn.Module):
    def __init__(self):
        super(Model_TL_A, self).__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.ModuleList()
        current_dim = 784
        for n_hidden in (300, 100, 50, 50, 50):
            self.layers.append(nn.Linear(current_dim, n_hidden))
            self.layers.append(nn.SELU())
            current_dim = n_hidden
        self.linear1 = nn.Linear(current_dim, 8)
            
    def forward(self, X):
        out = self.flatten(X)
        for layer in self.layers[:-1]:
            out = layer(out)
        out = self.linear1(out)
        return out

model_tl_a = Model_TL_A()

from torchsummary import summary
summary(model_tl_a, (1,28,28))

Here is my other model

class Model_TL_B_on_A(nn.Module):
    def __init__(self, Model_TL_A):
        super(Model_TL_B_on_A, self).__init__()
        self.model_tl_a = Model_TL_A()
#         modules = list(self.model_tl_a.children())[:-1]
#         self.model = nn.Sequential(*modules)

            
    def forward(self, X):
        out = self.model_tl_a(X)
#         out = self.model(X)
        return out

It works fine (since I have commented the lines that remove the last layer.

model_tl_b_on_a = Model_TL_B_on_A(Model_TL_A)

from torchsummary import summary
summary(model_tl_b_on_a, (1,28,28))

However, if I uncomment the two lines in Model_TL_B_on_A, and the last line in forward, I get the following error

TypeError: forward() takes 1 positional argument but 2 were given

Can someone help me out?

Thanks

2

Is the error when you define the model or when you try and get the summary?

3

Creating model works fine.

When I get the summary or train the model, it gives that error

5

I think it has something to do with the way you call modulelist. I am not 100% sure how that works. You can go to this forum post and try the solution there.

1 Like
6

The ModuleList works perfectly. It is evident by the correct training of Model_TL_A. It is just that reusing its layers causes problem. But I fixed it by modifying the code. No ModuleList - directly used an empty layer list. Works fine now. Thanks.

7

A easier way to transfer your learning is simply to load you model state dict and then modify the last layer instead of creating a new class.
do torch.save(model.state_dict(), "model_a.pth’) once the Model_TL_A has finished learning.
Then when you want to start training your model_B:
model_B = Model_TL_A()
model_B.load_state_dict(torch.load(“model_a.pth”))
model.linear1 = nn.Linear(in_features=784, out_features=number_of_outputs)

Of course there are variants, but this is most simple way to go in my opinion.

1 Like