做法其实很简单,只需要套用原来的多进程方法,比如异步map()方法,还可以在进程结束后统一收集训练的模型。
但是可能会遇见optimizer.step()报错,一般问题都是因为安装的torch版本是GPU版本,只需要将版本改为CPU版本即可。因为GPU版本的torch会在BP时检查cuda的健康情况,而python自带的multiprocessing库无法调用cuda。
如果想在多进程的情况下用GPU,则需要使用torch.multiprocessing库,这个库是python自带multiprocessing库的封装,继承了大部分功能,并在其基础上增加了多机多卡等的训练模式。