I am doing metric learning and a single iteration of the training loop looks as follows:
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()

Or in words: I am performing batched operations (CNN) in learnable_network_1 and then unbatched pairwise operations in learnable_network_2 as loss_fn. 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 backward per single optim.step(). 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.
One way to address that would be to just call backward on each loss and not just at total_loss at the end, however that would mean backpropagating through learnable_network_1 quadratically many times.
An alternative, which I don’t know how to implement, is to backpropagate in two steps: in the loop, I would call some backpropagate_until_batchwise_descriptors(loss) which treats batchwise_descriptors 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 backpropagate_through_cnn(batchwise_descriptors.grad) to find CNN gradients. This approach would mean I only need to store a single pseudo_distances matrix in memory at any given time, greatly reducing the memory strain of this operation.
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?