Hi I’m trying to do MAML with learn2learn but it occurred to me that the loss for the meta-model is extremely expensive in terms of memory because it includes the computational graph of each base learner for each task inside a mini batch which is completely unnecessary.
What I’m trying to do is just accumulate the weight updates from each base learner (+ task pair) without the entire computational graph. How can you do this without actually applying the weight updates?
For example is it possible/easy to store a copy of the model where you zero all parameters then apply the weight updates from the real model to that model one task at a time. Then once you are done add that models parameters back into the original model?
I’m particularly interested in if anyone knows of a better solution?