openai / supervised-reptile

Code for the paper "On First-Order Meta-Learning Algorithms"
https://arxiv.org/abs/1803.02999
MIT License
989 stars 210 forks source link

Parallel version #11

Open bkj opened 6 years ago

bkj commented 6 years ago

Any plans to release a multi-GPU version of this? It looks like we should be able to run the meta_batch_size iterations of the outer loop in reptile.train_step in parallel on separate GPUs.

(I may have a shot at implementing it if there are no plans ATM, and if you think it'd give a nontrivial speedup.)

~ Ben

unixpickle commented 6 years ago

Reptile is definitely simple to scale across multiple machines, since each machine just has to run a separate inner loop and then average the parameters at the end. One thing about this implementation that makes it a bit tricky to parallelize is that Adam's parameters are updated sequentially for each task in the meta-batch. That is almost certainly unimportant, though, and there's other reasonable ways to do it.