Backpropagating in two steps

jatentaki: <p>I am doing metric learning and a single iteration of the training loop looks as follows:</p>
<pre><code class="lang-auto">batch = get_new_batch()
batchwise_descriptors = learnable_network_1(batch)

total_loss = 0
for i in range(batch.shape[0]):
    for j in range(i+1, batch.shape[0]):
        pseudo_distances = learnable_network_2(
              batchwise_descriptors[i],
              batchwise_descriptors[j]
         )
        loss = loss_fn(pseudo_distances)
        total_loss += loss

total_loss.backward()
optim.step()
</code></pre>
<p>Or in words: I am performing batched operations (CNN) in <code>learnable_network_1</code> and then unbatched pairwise operations in <code>learnable_network_2</code> as <code>loss_fn</code>. This unbatched step is the memory bottleneck because the pseudo distance matrices are very big, so I end up having to run with batch size of 2 and run multiple steps of <code>backward</code> per single <code>optim.step()</code>. This is highly non-optimal because in general my gradient estimates benefit from the quadratic behavior of the pairwise loop: the bigger batch, the better gradient estimates per single run of the CNN.</p>
<p>One way to address that would be to just call <code>backward</code> on each <code>loss</code> and not just at <code>total_loss</code> at the end, however that would mean backpropagating through <code>learnable_network_1</code> quadratically many times.</p>
<p>An alternative, which I don’t know how to implement, is to backpropagate in two steps: in the loop, I would call some <code>backpropagate_until_batchwise_descriptors(loss)</code> which treats <code>batchwise_descriptors</code> as leaf variables and just accumulates the gradient in them (without going into the CNN) and then, after the pairwise iteration is over, I’d call <code>backpropagate_through_cnn(batchwise_descriptors.grad)</code> to find CNN gradients. This approach would mean I only need to store a single <code>pseudo_distances</code> matrix in memory at any given time, greatly reducing the memory strain of this operation.</p>
<p>Unless I’m very mistaken this two-step backpropagation is mathematically correct and I think it should be possible to implement, though I have no idea how difficult it is and at which point I’d need to integrate with autograd. Could I ask for some input on that?</p>

albanD: <p>Hi,</p>
<p>My first question would be why can’t you run the second network in a batched manner? You could create all the combinations of batch by doing something like: <code>batchwise_descriptors.unsqueeze(0).expand(full_size)</code> for the first one and <code>batchwise_descriptors.unsqueeze(1).expand(full_size)</code> for the second one. Where full_size is <code>(batch_size, batch_size, your_other_dims...)</code>.<br>
That would definitely be the most efficient in terms of speed.</p>
<p>Not if you have hard constraints that you cannot do that (or memory issues), you can do the following:<br>
(I haven’t tested the code so there might be a typo, but this is the spirit).</p>
<pre><code class="lang-python">batch = get_new_batch()
optim.zero_grad() # Make sure all .grads are 0
batchwise_descriptors = learnable_network_1(batch)

# Create a new leaf Tensor that will accumulate all the grads
# coming from net 2
net_2_in = batchwise_descriptors.detach().requires_grad_()

total_loss = 0
for i in range(batch.shape[0]):
    for j in range(i+1, batch.shape[0]):
        pseudo_distances = learnable_network_2(
              net_2_in[i],
              net_2_in[j]
         )
        local_loss = loss_fn(pseudo_distances)
        local_loss.backward() # Accumulate in in learnable_network_2 and net_2_in gradients
        total_loss += loss.detach() # Just used for bookeeping, no gradient should flow back

# Now backwarprop the whole thing in net 1
batchwise_descriptors.backward(net_2_in.grad)
optim.step()
</code></pre>

jatentaki: <p>Great, it looks much simpler than I expected <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Why can’t I do it in a batched way? That’s because I hid a lot of complexity in this example, in reality I extract a random number of descriptors from each batch element so the pseudo-distance matrices are of different shapes from each other (which would require a huge amount of padding) and even if that worked, the <code>loss_fn</code> uses external, non-batchable operations.</p>
<p>Thanks a lot!</p>

albanD: <p>Ok sounds good!<br>
Then the second example should do what you want <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Note that you can play with the tradeof between memory and speed by accumulating a bunch of losses in the double for loop before calling backward. Doing so will use more memory but will speed things up a little bit.<br>
Like for example accumulate the loss for the <code>j</code> loop and call backward once at the end.</p>