首先,我们先把我们的pytorch版本提升到2.0.1,这样会防止很多不必要的报错(但这样的坏处就是我们没有办法使用nvidia-apex进行加速了,除非等到版本的更新
DistributedSampler
我们使用分布数据采样器来保证我们的每个进程(每一块卡)都接受不同的输入
这就是并行计算中的数据并行
对DataLoader的设置
在sampler参数中加入DistributedSampler(dataset),然后把shuffle调为False
有关Optimizer
由于我们的每个进程接受的不同的输入,所以产生的损失函数也不一样,反向传播操作中的梯度也是不一样的,所以这会导致我们每个设备上的模型有不同的参数,所以这里DDP(DistributedDataParallel)会采用一个桶环状的AllReduce算法,这使得梯度计算和不同进程的交流重叠起来,这样就保证了我们每个模型的复制品都有相同的梯度。
多GPU训练
代码(必要的几个package)
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group , destroy_process_group
import osdef ddp_setup(rank , world_size):"""Args:rank : Unique identifier of each processworld_size : Total number of processes"""os.environ["MASTER_ADDR"] = "localhost"os.environ["MASTER_PORT"] = "12355" # can be any free port on you machinetorch.cuda.set_device(rank) # this is importantinit_process_group(backend="nccl",rank=rank,world_size=world_size)def Training_function(rank:int,world_size:int,epoch,batchsize....):ddp_setup(rank,world_size)"""your codenet = DDP(net.cuda(),device_id=[rank]"""destroy_process_group()if __name__ == "__main__":world_size = torch.cuda.device_count()mp.spawn(Training_function,args=(world_size:int,epoch,batchsize....),nprocs=world_size)"""这里我们不需要传递rank因为mp.spawn会自动传递"""