在我们训练网络的时候,会出现GPU显存占用不满的情况
想最大化利用gpu加速训练可以从以下两个方向入手:
- 使用多进程分别训练权重网络,然后再平均权重网络
- 使用多进程获取环境参数,再交给主进程更新权重网络(A3C算法框架)
为什么是多进程(Multiprocessing)
python底层是基于C的伪多线程,说是多线程实际上是挂了一个全局锁的单线程。
加了一些线程切换反而会变慢,这时候需要用多进程去加速
多进程使用上和多线程差不多
import multiprocessing as mp
mp1 = mp.Process(target=某某任务, args=(参数1,))
mp1.start()
mp1.join()
进程间传输参数
multiprocessing .Pipe
提供了双端管道功能,可以互传参数
(con1, con2) = Pipe()
#通过发:
con1.send(“发送’)
#以及收:
con2(pipe.recv())
通过这个功能 就可以实现上述两种功能
或者使用model.share_memory()功能,将model传入worker中 然后load_state_dict即可,这块有兴趣可以找我要,实际上和上述实现方法效果一样
代码实现
所谓的A3C
就是:采样过程并行、训练过程并行
这里的实现只取前者
#建立管道连接
pipe_dict = dict((i, (pipe1, pipe2)) for i in range(processes) for pipe1, pipe2 in (mp.Pipe(),))
child_process_list = []
for i in range(processes):
#传入管道参数
pro = mp.Process(target=processEpoch, args=(pipe_dict[i][1],))
child_process_list.append(pro)
#启动进程
[p.start() for p in child_process_list]
然后就可以利用多进程去跑这个processEpoch
里面实际上就是跑环境,然后获取reward next_state
然后全部打包给主进程
一般正向传播过程建议只用CPU,除非你的GPU非常充裕。
如果你的GPU非常强大,可以参考文章最后的方法
实战代码
因为是从程序中抽取出来的,可能比较乱,主需要看核心数据抽取即可:
(其实就是多个进程玩游戏,然后把环境啥的发给大脑训练)
ctx = mp.get_context("spawn")
config = configparser.ConfigParser()
config.read("hyperConfig.conf", encoding="utf-8")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 超参数
BATCH_SIZE = config.getint('HYPERPARA','BATCH_SIZE')
LR = config.getfloat('HYPERPARA','LR') # learning rate
EPSILON = config.getfloat('HYPERPARA'