How to speed up training about greedy_infomax HOT 1 CLOSED

One way to reduce memory usage could be to train the modules individually, i.e. one after another. While the overall procedure for this is not implemented in this codebase, you could start by setting --train_module 0 to train the first module. After that, it would be advisable to save the outputs of this model as a new dataset and train the next module on it.

If you have several GPUs available, you could potentially parallelize this process as well by training one module per GPU and by saving the outputs of each module as the inputs for the next module in regular intervals.

Here's a rough schematic of what I mean, where each of the gray boxes would represent a separate GPU:

Hope this helps!

from greedy_infomax.