Reuse existing model

username: mmg
cooked: <p>I am trying to reuse a model (transfer learning). I want to remove the last layer (and will later add another layer)</p>
<p>My existing model is</p>
<pre><code class="lang-auto"> 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
</code></pre>
<pre><code class="lang-auto">model_tl_a = Model_TL_A()

from torchsummary import summary
summary(model_tl_a, (1,28,28))
</code></pre>
<p>Here is my other model</p>
<pre><code class="lang-auto">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
</code></pre>
<p>It works fine (since I have commented the lines that remove the last layer.</p>
<pre><code class="lang-auto">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))
</code></pre>
<p>However, if I uncomment the two lines in Model_TL_B_on_A, and the last line in <code>forward</code>, I get the following error</p>
<pre><code class="lang-auto">TypeError: forward() takes 1 positional argument but 2 were given
</code></pre>
<p>Can someone help me out?</p>
<p>Thanks</p>

username: Dwight_Foster
cooked: <p>Is the error when you define the model or when you try and get the summary?</p>

username: mmg
cooked: <p>Creating model works fine.</p>
<p>When I get the summary or train the model, it gives that error</p>

username: Dwight_Foster
cooked: <p>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 <a href="https://discuss.pytorch.org/t/typeerror-forward-takes-1-positional-argument-but-2-were-given/104776/2">this</a> forum post and try the solution there.</p>

username: mmg
cooked: <p>The ModuleList works perfectly. It is evident by the correct training of <code>Model_TL_A</code>. 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.</p>

username: aguennecjacq
cooked: <p>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.<br>
do torch.save(model.state_dict(), "model_a.pth’) once the Model_TL_A has finished learning.<br>
Then when you want to start training your model_B:<br>
model_B = Model_TL_A()<br>
model_B.load_state_dict(torch.load(“model_a.pth”))<br>
model.linear1 = nn.Linear(in_features=784, out_features=number_of_outputs)</p>
<p>Of course there are variants, but this is most simple way to go in my opinion.</p>