torch.multiprocessing

torch.multiprocessing是具有额外功能的multiprocessing，其 API 与multiprocessing完全兼容，因此我们可以将其用作直接替代品。

multiprocessing支持 3 种进程启动方法：fork（Unix 上默认）、spawn（Windows 和 MacOS 上默认）和forkserver。要在子进程中使用 CUDA，必须使用forkserver或spawn。启动方法应该通过set_start_method()在if name == 'main’主模块的子句中使用来设置一次：

import torch.multiprocessing as mp

if __name__ == '__main__':
    mp.set_start_method('forkserver')
    ...

张量共享

import torch.multiprocessing as mp
import time

mat = torch.randn((200, 200))
print(mat.is_shared())

queue = mp.Queue()
q.put(a)
print(a.is_shared())
#False
#True

import torch
import torch.multiprocessing as mp
import time

def foo(q):
    q.put(torch.randn(20, 20))
    q.put(torch.randn(10, 10))
    time.sleep(3)

def bar(q):
    t1 = q.get()
    print(f"Received {t1.size()}")
    time.sleep(1) # 注意这里不能等待超过3，会导致foo结束，无法获取
    t2 = q.get()
    print(f"Received {t2.size()}")

if __name__ == "__main__":
    mp.set_start_method('spawn')

    queue = mp.Queue()
    p1 = mp.Process(target=foo, args=(queue,))
    p2 = mp.Process(target=bar, args=(queue,))

    p1.start()
    p2.start()

    p1.join()
    p2.join()

仅在Python 3中使用spawn或forkserver启动方法才支持在进程之间共享CUDA张量。

torch.multiprocessing.spawn

torch.multiprocessing.spawn(
    fn,
    args=(),
    nprocs=1,
    join=True, # join (bool) – 执行一个阻塞的join对于所有进程.
    daemon=False, # daemon (bool) – 派生进程守护进程标志。如果设置为True，将创建守护进程.
    start_method="spawn",
)

其中，fn 是要在子进程中运行的函数，args 是传递给该函数的参数，nprocs 是要启动的进程数。当 nprocs 大于 1 时，会创建多个子进程，并在每个子进程中调用 fn 函数，每个子进程都会使用不同的进程 ID 进行标识。当 nprocs 等于 1 时，会在当前进程中直接调用 fn 函数，而不会创建新的子进程。

join=true时，主进程等待所有子进程完成执行并退出，然后继续执行后续的代码。在这个过程中，主进程会被阻塞，也就是说，主进程会暂停执行，直到所有子进程都完成了它们的任务。

torch.multiprocessing.spawn 函数会自动将数据分布到各个进程中，并在所有进程执行完成后进行同步，以确保各个进程之间的数据一致性。同时，该函数还支持多种进程间通信方式，如共享内存（Shared Memory）、管道（Pipe）等，可以根据具体的需求进行选择。

比如下面的fn

def main_worker(gpu, args): # gpu参数控制进程号
	args.rank = gpu # 用rank记录进程id号
	dist.init_process_group(backend='nccl', init_method=args.dist_url, world_size=args.num_gpus,
	                                rank=args.rank)
	torch.cuda.set_device(gpu)  # 设置默认GPU  最好方法哦init之后，这样你使用.cuda()，数据就是去指定的gpu上了
	
	# 定义模型， 转同步BN
	model = xxx
	model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
	model.cuda()
	model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu],
	                                                      find_unused_parameters=True )
	#定义数据集
	train_dataset = xxxx
	# 注意这一步，和单卡的唯一区别。这个sample能保证多个进程不会取重复的数据。shuffle必须设置为False（默认）
	train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
	train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=args.batch_size,
	                   num_workers=args.workers, pin_memory=True, sampler=train_sampler)
	...
	
	if args.rank == 0 ： # 主进程
    torch.save(xx)
    print(log)

启动代码

import torch.multiprocessing as mp
import torch.distributed as dist

if __name__ == '__main__':
	# import configuration file
	# load json or yaml, argsparse
	args = xxxxx
    # 接下来是设置多进程启动的代码
    # 1.首先设置端口，采用随机的办法，被占用的概率几乎很低.
    port_id = 29999
    args.dist_url = 'tcp://127.0.0.1:' + str(port_id)
    # 2. 然后统计能使用的GPU，决定我们要开几个进程,也被称为world size
    args.num_gpus = torch.cuda.device_count()
    # 3. 多进程的启动
   	torch.multiprocessing.set_start_method('spawn')
    mp.spawn(main_worker, nprocs=args.num_gpus, args=(args,))

multiprocessing.Pool与torch.multiprocessing.Pool

multiprocessing.Pool创建一个进程池，每个进程都分配一个单独的内存空间。它是一个上下文管理器，因此可以在语句中使用with

with mp.Pool(processes=num_workers) as pool:

等价于

pool = mp.Pool(processes=num_workers)
# do something
pool.close()
pool.join()

阻塞

import time
import multiprocessing as mp

def foo(x, y):
    time.sleep(3)
    return x + y

with mp.Pool(processes=4) as pool:
    a = pool.apply(foo, (1, 2))
    b = pool.apply(foo, (3, 4))
    print(a, b)
# 3 7
#---
#Runtime: 6.0 seconds

创建一个包含 4 个工作进程的池，然后向池中提交两个任务来运行。由于apply是阻塞调用，因此主进程将等到第一个任务完成后再提交第二个任务。这基本上是无用的，因为这里没有实现并行性。

with mp.Pool(processes=4) as pool:
    handle1 = pool.apply_async(foo, (1, 2))
    handle2 = pool.apply_async(foo, (3, 4))

    a = handle1.get()
    b = handle2.get()

    print(a, b)
# 3 7
#---
#Runtime: 3.0 seconds

非阻塞

apply_async是非阻塞的并AsyncResult立即返回一个对象。然后我们可以使用它get来获取任务的结果。

注意get会阻塞直到任务完成；apply(fn, args, kwargs)相当于apply_async(fn, args, kwargs).get().

还可以加回调函数

def callback(result):
    print(f"Got result: {result}")

with mp.Pool(processes=4) as pool:
    handle1 = pool.apply_async(foo, (1, 2), callback=callback)
    handle2 = pool.apply_async(foo, (3, 4), callback=callback)
#Got result: 3
#Got result: 7
#---
#Runtime: 3.0 seconds

map

map将输入的可迭代对象划分为块，并将每个块作为单独的任务提交到池中。然后收集任务的结果并以列表的形式返回。

阻塞

import multiprocessing as mp
import time

def foo(x):
    print(f"Starting foo({x})")
    time.sleep(2)
    return x

with mp.Pool(processes=2) as pool:
    result = pool.map(foo, range(10), chunksize=None)
    print(result)

Starting foo(0)
Starting foo(2)
Starting foo(1)
Starting foo(3)
Starting foo(4)
Starting foo(6)
Starting foo(5)
Starting foo(7)
Starting foo(8)
Starting foo(9)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
---
Runtime: 12.0 seconds

在这种情况下，chunksize 自动计算为 2。这意味着可迭代对象被分为 5 个大小为 2: 的块[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]。

map是一个阻塞调用，因此它将等待所有任务完成后再返回。与 apply类似

首先，前两个块[0, 1], [2, 3]分别被提交给2个worker，此时worker0执行0worker1执行2，然后worker0执行1worker1执行3；

然后接下来的两个块[4, 5], [6, 7]被提交。最后，最后一个块[8, 9]被提交给任一worker。

三次提交，每次提交运行2*2秒，共计12s，这不是最优的。在这种情况下，如果我们显式地将 chunksize 设置为 1 或 5，则运行时间将为 10 秒，这已经是最好的了。

非阻塞

with mp.Pool(processes=2) as pool:
    handle = pool.map_async(foo, range(10), chunksize=None)
    # do something else
    result = handle.get()
    print(result)

starmap

map只是将可迭代的元素传递给函数。如果想应用一个多参数函数，我们要么必须传入一个列表并将其解压到函数内。

使用starmap. 对于可迭代的每个元素，starmap将其解压到函数的参数中。

def bar(x, y):
    print(f"Starting bar({x}, {y})")
    time.sleep(2)
    return x + y

with mp.Pool(processes=2) as pool:
    pool.starmap(bar, [(1, 2), (3, 4), (5, 6)])

starmap是同步的，异步用starmap_async

https://tokudayo.github.io/multiprocessing-in-python-and-torch/#torchmultiprocessing
https://docs.python.org/zh-cn/3.7/library/multiprocessing.html