Multiple Inner Loops

[pytorch]

Multiple Inner Loops

[Chillee]

get_loss_for_task(x1, y1, x2, y2) gets you the loss for a single inner task (i.e. a single "inner loop", although it's only doing one step of the inner training loop here). This single inner task is (x1, y1, x2, y2), and consists of 2 examples pairs. So, let's say that x1 has shape[5].

So, if you vmap over get_loss_for_task, then you'd now pass in [B, 5] for x1 instead of just [5]. This corresponds to doing 5 runs of get_loss_for_task, or running 5 "inner training loops" in parallel.

[pharringtonp19]

@Chillee Thanks for the response, I should have been clearer in my original question.

As you note, "it's only doing one step of the inner training loop here". How would I adjust the code so that it does k steps instead of 1?

[Chillee]

https://github.com/zou3519/functorch/blob/main/examples/maml_regression/evjang_transforms.py#L73

Instead of doing only one gradient step here, you need to pass in multiple examples and do multiple gradient steps in a loop.

[pharringtonp19]

@Chillee Thanks