Multiple model.forward / Multiple models.forward followed by one loss.backward

Andre_Godinho: <p>Hello everyone,</p>
<p>I want to know the best implementation out of three similar implementations regarding training a bi-encoder model in PyTorch with NLL (as a triplet loss) in terms of less memory allocation. Moreover, there are anchor, positive and negative outputs for the NLL and the differences among the implementations are:</p>
<ol>
<li>uses only one model and do a forward pass in each batch iteration for positives and negatives</li>
<li>uses two models (with two optimizers) and do a forward pass in each batch iteration for positives and negatives</li>
<li>uses one model and do a forward pass over all anchors, positives and negatives</li>
</ol>
<p>My <strong>goal</strong> is to <strong>reduce the allocated memory with positives and negatives in order to fit more negatives in the GPU</strong> for each batch.  I am not sure if the two first ones allocate less memory because they might be creating a separate computational graph for each forward pass. Also, I must do a forward pass over all anchors at once and call <code>loss.backward(retain_graph=True)</code> over batches of anchors, this is common in all implementations.</p>
<p>The problem with the third one is that I compute all the positives and negatives at once, and they are occupying space in memory.</p>
<ol>
<li>Implementation 1 (pseudocode)</li>
</ol>
<pre><code class="lang-auto">optimizer.zero_grad()

all_anchors_embed = bi_encoder.forward(all_anchors_input)

for b_anchors_embed, b_pos_inputs, b_neg_inputs in zip(all_anchors_embed, all_pos_input, all_neg_input):
    b_pos_embed, b_neg_embed = bi_encoder.forward(b_pos_inputs, b_neg_inputs)

    loss = NLL(b_anchors_embed, b_pos_embed, b_neg_embed)
    loss.backward(retain_graph=True)

optimizer.step()
</code></pre>
<ol start="2">
<li>Implementation 2: (pseudocode)</li>
</ol>
<pre><code class="lang-auto">optimizer_encoder_1.zero_grad()
optimizer_encoder_2.zero_grad()

all_anchors_embed = encoder_1.forward(all_anchors_input)

for b_anchors_embed, b_pos_inputs, b_neg_inputs in zip(all_anchors_embed, all_pos_input, all_neg_input):
    b_pos_embed, b_neg_embed = encoder_2.forward(b_pos_inputs, b_neg_inputs)

    loss = NLL(b_anchors_embed, b_pos_embed, b_neg_embed)
    loss.backward(retain_graph=True)

optimizer_encoder_1.step()
optimizer_encoder_2.step()
</code></pre>
<ol start="3">
<li>Implementation 3 (similar to <a href="https://discuss.pytorch.org/t/multiple-model-forward-followed-by-one-loss-backward/20868">Multiple model.forward followed by one loss.backward</a>): (pseudocode)</li>
</ol>
<pre><code class="lang-auto">optimizer.zero_grad()

all_anchors_embed, all_pos_embed, all_neg_embed = bi_encoder.forward(all_anchors_input, all_pos_input, all_neg_input)

for b_anchors_embed, b_pos_embed, b_neg_embed in zip(all_anchors_embed, all_pos_embed, all_neg_embed):
    loss = NLL(b_anchors_embed, b_pos_embed, b_neg_embed)
    loss.backward(retain_graph=True)

optimizer.step()
</code></pre>
<p>Which one do you think will fit more negatives in memory?</p>
<p>Take in mind that the bi-encoder in 1) and 3) will have two encoders inside, one to encode the anchors and another one to encode the positives and negatives.</p>