多进程/多GPU推理，提速5倍以上

最近项目用到SuperGlue匹配SuperPoint特征点，实验发现，对每个img-img对进行SuperGlue推理，一个img_pair需要~1s，而项目中所有类的img_pair数量~60*5K个，时间接近42小时，如果数量更多，单匹配花费时间就大于2天，难以等待。

看推理只占用了一张GPU卡，而且只用了1.5G显存，看服务器有多GPU卡空闲，想通过多进程并行推理，性能加速。

一、多进程

import multiprocessing as mp # 这里不需要使用pytorch的multiprocessing
mp = mp.get_context('spawn')
pool = []
num_process = 4  # 开4进程
total_num = len(target_loc)   # total pairs num
phrase = total_num // num_process + 1

for i in range(num_process):
    divied_target = target_loc[i*phrase:(i+1)*phrase]   # 所有的数据分成4份

   process = mp.Process(target=multi_process_pred,  args=(divied_target, i))   # 这里可以把每个进程需要的参数传进去
    # 每个进程都执行multi_process_pred，在multi_process_pred中包含load 深度学习model，根据img_id - img_id读取图像的局部特征，通过superglue类，推理图像之间的匹配信息，保存匹配结果到对应这个进程的文件。
    process.start()
    pool.append(process)

for p in pool:
    p.join()  # 等待所有进程执行完毕

在执行main函数的shell脚本中通过CUDA_VISIBLE_DEVICES=3指定那块GPU，然后，开4进程会提速大于2倍。

二、多GPU、多进程

在执行main函数的shell脚本中通过CUDA_VISIBLE_DEVICES=1,2,3指定多块GPU。

通过torch.cuda.device_count()判断多少GPU可以使用。

平均将多个进程分配到每个GPU，同时对每个进程指定具体的gpu_idx，

通过self.device = 'cuda:'+str(gpu_idx) if torch.cuda.is_available() else 'cpu'

传到SuperGlue(config).eval().to(self.device)中的self.device，在每块GPU上推理。

三、解决CUDA out of memory. Tried to allocate 98.00 MiB

1.  https://github.com/pytorch/pytorch/issues/16417

with torch.no_grad():
    # Your eval code here

2. https://stackoverflow.com/questions/59129812/how-to-avoid-cuda-out-of-memory-in-pytorch

import torch
torch.cuda.empty_cache()

del pred
gc.collect()

torch.cuda.memory_summary(device=None, abbreviated=False)

总结：（1）开的进程少一些，每个卡留~2G显存的余地，应对load数据可能的上下波动。【重要】

（2）我加了上面所有方法控制显存增加，结果有效，显存没有增加。

四、多进程共享变量

在每个进程单独load所有图像的特征，这样内存占用比较大，特别是把整个数据库图像的特征一起load时。所以最好多进程共享load的特征。

import multiprocessing as mp # 这里不需要使用pytorch的multiprocessing                                    
mp = mp.get_context('spawn')
manager = mp.Manager()
my_share_var = manager.dict() # 这里存入所有图像的特征，如果是从h5读取的特征，要转换到array

然后，

pool = []
for idx in range(num_process):
    process = multiprocessing.Process(target=func, args=(my_share_var, other_vars))
    process.start()
    pool.append(process)
for proc in pool:
    proc.join()

实验比较：

1.load整个数据库所有图像的特征（12G），开2个GPU，4个进程：

 -- 共享变量时，占用内存：34G

 -- 每个进程单独load所有特征时，占用内存54G，因为快超了总内存，崩了一个进程，只起了3个进程。

2. 依次遍历每个子类，load每个子类图像的特征（1~2G），开4个GPU，19个进程：

 -- 共享变量时，占用内存：30G

 -- 每个进程单独load所有特征时，占用内存33G。（可能开的进程太多，共享变量在共享到多进程时，需要增加内存）

结论：（1）如果load的数据量特征比较大，使用共享内存。

（2）如果依次遍历每个子类，load每个子类图像的特征，不使用共享内存。

（3）多进程并行数量适量即可：

       --  1个GPU最多开5个进程；

       --   GPU:0最多开4个进程，因为其他GPU会在第0块GPU占用显存；

       --   4GPU，总进程不超过20个；

       --   不要超过服务器CPU总核数。

       实验测试：进程开的太多（>20个），运行效率反而下降。

感谢以下博客指点迷津：

https://blog.csdn.net/qq_34914551/article/details/110285386

https://zhuanlan.zhihu.com/p/166091204