Pytorch 多进程在单卡上测试

有些炼丹师可能机器不足，只有一张卡，然后训练完成了，想要测试的时候，受限于图像样本size不一致，不能合并到一个batch中。当然解决方案很多，但有一种更自然的办法，既使用任意分辨率的输入，同样使用多进程在单卡上并行执行数据的测试，从而加快测试速度，赶上dideline。

大家都知道Pytorch官方推广大家使用分布式多卡并行计算，其原理是每个进程都使用一个GPU。那我们现在要做的就是多进程都在一个GPU上，每个进程处理一批数据，从而加快处理速度。

前言

• 模型在GPU上，Pytorch规定：多进程的启动方法必须是"spawn"，所以使用map_async或者apply_async这类方法是不行了。

• 执行这套流程的大致思路是，我们可以先把要处理的数据加入到容器中（比如字典或者元祖列表），然后把容器中的数据分段，分段的数目就是进程数，然后每个进程处理自己对应的那些数据。由于输入的size不一样，所以我们要一张一张的处理，但因为开了多进程，实际的batchsize等于进程数目。

假设我们现在有一批数据，我们先要把这批数据的目录地址加到一个容器中，

    target_loc = make_dataset(os.path.join(img_root, training_subfix), True)
    target_loc.extend(make_dataset(os.path.join(img_root, val_subfix), False))
    target_loc.sort(key=lambda x: x[0]) # 保证一致
    print(len(target_loc))

make_dataset的作用就是把所有的数据的地址加入到target_loc中。target_loc是一个list。
然后开始多进程启动

import multiprocessing as mp # 这里不需要使用pytorch的multiprocessing
mp = mp.get_context('spawn')
pool = []
n_proc = 4  # 开4进程
total_num = len(target_loc)
phrase = total_num // n_proc + 1
for i in range(n_proc):
    divied_target = target_loc[i*phrase:(i+1)*phrase]  # 所有的数据分成4份
    process = mp.Process(target=prediction_by_dfsd, args=(divied_target,i)) 
    # 每个进程都执行prediction_by_dfsd
    process.start()
    pool.append(process)
for p in pool:
    p.join()  # 等待所有进程执行完毕

prediction_by_dfsd这个函数就是测试数据的函数，这个函数需要完成的是，从分段的容器中读取图像，然后一张张的预测，然后保存结果。
另外我需要说明的是，spawn的启动进程的方式和我们平常使用的map_async等不一样，从spawn启动的进程，是无法访问到主进程的任何变量的。所有我们需要在prediction_by_dfsd完成模型定义，导入参数等操作。当然一些变量可以通过Process的args传进去。
所以和使用分布式训练一样，每个进程都会定义一次模型。只不过我们这个方法，是在单卡上使用多进程，一样能使用并行预测，满足了在一些情况，比如样本size不一样，不能合并到batch中，而一张一张预测又太慢的问题。

最后

当然我后面反应过来了。只用使用dataloader，batchsize为1，一样可以使用标准的pytorch分布式流程使用多进程在单卡上的测试，只不过不需要设置gpu id就行了。所有的模型都往一个gpu上搬运。不过呢，对于简单的任务，没有必要再写一个dataloader，所以我这个办法还是有价值的。