Un breve resumen del entrenamiento distribuido de pytorch.

Esto es relativamente simple; vaya a otro artículo para obtener más detalles: https://blog.csdn.net/qq_36276587/article/details/123913384

Un breve resumen del uso de pytorch para la capacitación distribuida con múltiples tarjetas en una sola máquina, principalmente el uso de algunas API clave y el proceso de capacitación distribuida, pytorch versión 1.2.0 está disponible

 Inicialice el método de comunicación GPU (NCCL)
 

import torch.distributed as dist

torch.cuda.set_device(FLAGS.local_rank)
dist.init_process_group(backend='nccl')
device = torch.device("cuda", FLAGS.local_rank) #自己设置

 Carga de datos distribuidos

train_sampler = torch.utils.data.distributed.DistributedSampler(traindataset)
train_loader = torch.utils.data.DataLoader(
        traindataset, batch_size=batchSize,
        sampler=train_sampler,
        num_workers=4, pin_memory=True,#drop_last=False,
        collate_fn=alignCollate(imgH=imgH, imgW=imgW, keep_ratio=FLAGS.keep_ratio))
#pytorch的DataLoader格式处理训练标签

Modelo de entrenamiento distribuido

#初始化后的模型使用分布式训练
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)  ## 同步bn
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[FLAGS.local_rank],
                                                      output_device=FLAGS.local_rank)

empezar a entrenar

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 train_distributed.py