Modifying Weight of Pretrained Model in `forward()` Makes Training Slowing

malioboro: <p>I have two parameters, A and B, that I need to put to replace all the weight of the pre-trained model. So I want to utilize the forward calculation of the model but not the weight.</p>
<p>I want to modify the weight of the model W = A + B, where A is a fixed tensor (not trainable), but B is a trainable parameter. So, in the end, my aim is to train B in the structure of the pre-trained model.</p>
<p>This is my attempt:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    
    def __init__(self, pre_model, B):
    
        super(Net, self).__init__()
        self.B = B
        with torch.no_grad():
            self.model = copy.deepcopy(pre_model)
            for params in self.model.parameters(): 
                params.requires_grad = False
    
    def forward(self, x, A):
        for i, params in enumerate(self.model.parameters()):
            params.copy_(A[i].detach().clone())
            params.add_(self.B[i])    
            params.retain_grad()

        x = self.model(x)
        return x
</code></pre>
<p>I checked in the process, B was already trained. But the problem is in every iteration the training process keeps getting slower:</p>
<p>Epoch 1:<br>
24%|██▍       | 47/196 [00:05&lt;00:23,  6.44it/s]<br>
57%|█████▋    | 111/196 [00:18&lt;00:19,  4.28it/s]<br>
96%|█████████▋| 189/196 [00:41&lt;00:02,  2.90it/s]<br>
Epoch 2:<br>
6%|▌         | 11/196 [00:04&lt;01:14,  2.50it/s]</p>
<p>I think I have detached all the parameters correctly, but I am not sure why it’s happened. I hope somebody here can help me to find out what exactly happened here.<br>
Thanks</p>

ptrblck: <aside class="quote no-group" data-username="malioboro" data-post="1" data-topic="169602">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/malioboro/48/1660_2.png" class="avatar"> malioboro:</div>
<blockquote>
<p>where A is a fixed tensor (not trainable), but B is a trainable parameter.</p>
</blockquote>
</aside>
<p>Based on your code snippet you are detaching <code>A</code>, which is the fixed tensor, while you are adding <code>B</code> to <code>params</code> potentially including its entire computation graph. Could you double check this, please?</p>

malioboro: <p>Basically, this is what I did to get A and B (it’s a little bit different, but I’ve tested the code below, and it’s still giving slowing iteration):</p>
<pre><code class="lang-auto">b = []
A = []
for params in list(pre_model.parameters()):
    A.append(torch.rand_like(params))
    b_temp = nn.Parameter(torch.rand_like(params))
    b.append(b_temp.detach().clone())
B = nn.ParameterList(b)

modelwithAB = Net(pre_model, B)
# ...
# in the training iteration
out = modelwithAB(image, A)
</code></pre>

malioboro: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> , from <a href="https://discuss.pytorch.org/t/how-to-modify-weights-of-a-layer-before-the-layer-is-applied-to-input/87096/4">another discussion</a> I read that you suggest using <code>no_grad()</code> context when modifying model parameters. But, in my case, here I didn’t use it because I want B to be updated by the optimizer through <code>add_</code> operation.</p>
<p>I tried to use <code>no_grad()</code>, and it made the training time more stable (not slowing), but as I expected, it made B not update. I am not sure which one is close to the solution.</p>

ptrblck: <p>I don’t know how exactly your model is used and where the computation graph might be stored as I cannot reproduce any increase in memory usage using:</p>
<pre><code class="lang-python">
device = 'cuda'
pre_model = models.resnet18().to(device)
b = []
A = []
for params in list(pre_model.parameters()):
    A.append(torch.rand_like(params))
    b_temp = nn.Parameter(torch.rand_like(params))
    b.append(b_temp.detach().clone())
B = nn.ParameterList(b)

modelwithAB = Net(pre_model, B)
optimizer = torch.optim.Adam(modelwithAB.parameters(), lr=1e-3)

image = torch.randn(2, 3, 224, 224).to(device)
print(torch.cuda.memory_allocated()/1024**2)

for _ in range(10):
    optimizer.zero_grad()
    out = modelwithAB(image, A)
    out.mean().backward()
    optimizer.step()
    print(torch.cuda.memory_allocated()/1024**2)
</code></pre>
<p>The <code>print</code> statement shows an approx. constant memory usage, which looks correct.</p>

malioboro: <p>Thank you, <a class="mention" href="/u/ptrblck">@ptrblck</a>, for trying to reproduce this error.</p>
<p>I have modified your code by adding <code>tqdm</code> and logging the time for <code>backward()</code> and <code>step()</code>, and even though the memory usage is still the same, you can see that the training time is slowing down.</p>
<pre><code class="lang-python">for i in tqdm(range(300)): # longer run, added tqdm
    optimizer.zero_grad()
    out = modelwithAB(image, A)
    start = time.time() # for logging the backpropagation's duration
    out.mean().backward()
    optimizer.step()
    if i%40==0:
        print("-", torch.cuda.memory_allocated()/1024**2, "-", time.time()-start)
</code></pre>
<p>the result from google colab (<a href="https://colab.research.google.com/drive/1L16Rs3rlNrtY-vJ0q4uz-40a-tJCggur?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a>):</p>
<pre><code class="lang-python">  1%|          | 3/300 [00:00&lt;00:10, 28.02it/s] - 1457.7119140625 - 0.02620530128479004
 14%|█▍        | 42/300 [00:02&lt;00:18, 14.10it/s] - 1365.5771484375 - 0.06569838523864746
 27%|██▋       | 82/300 [00:06&lt;00:29,  7.47it/s] - 1365.5693359375 - 0.12723588943481445
 41%|████      | 122/300 [00:13&lt;00:33,  5.37it/s] - 1365.5693359375 - 0.17061519622802734
 54%|█████▎    | 161/300 [00:21&lt;00:33,  4.10it/s] - 1365.5693359375 - 0.23227190971374512
 67%|██████▋   | 201/300 [00:32&lt;00:30,  3.30it/s] - 1365.5693359375 - 0.29410719871520996
 80%|████████  | 241/300 [00:46&lt;00:20,  2.82it/s] - 1365.5693359375 - 0.3430461883544922
 94%|█████████▎| 281/300 [01:02&lt;00:07,  2.41it/s] - 1365.5693359375 - 0.40257716178894043
100%|██████████| 300/300 [01:10&lt;00:00,  4.26it/s]
</code></pre>
<p>I have tried to remove <code>tqdm</code>, but the training time is still slowing down.</p>

malioboro: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> I am sorry keep mentioning you, but can you help me check the code above I still got stuck in this problem, and I don’t have any clue how to track/debug in this situation</p>