ddp Multi-Card-Trainingsbrenner-Rekord

Vorwort

Ich habe zuvor den Open-Source-Code anderer Leute verwendet, um ihn zu ändern, aber vor kurzem muss ich ihn von Grund auf neu schreiben. Das Experiment erfordert ein Training mit mehreren Karten, daher werde ich es aufzeichnen.

Grundgerüst

from torch.utils.data.distributed import DistributedSampler

 # 1) 初始化
 # torch.distributed.init_process_group(backend="nccl", init_method='env://', rank=0, world_size=torch.cuda.device_count())
  torch.distributed.init_process_group(backend="nccl")
  # 2) 配置每个进程的gpu
  local_rank = torch.distributed.get_rank()
  torch.cuda.set_device(local_rank)

  device = torch.device("cuda", local_rank)
  # 3) 使用DistributedSampler
  trainloader = DataLoader(trainset, batch_size=opt.batch_size, pin_memory=True,
                          num_workers=opt.num_workers, sampler=DistributedSampler(trainset, shuffle=True))

# 4) 封装之前要把模型移到对应的gpu
if num_gpus >= 1:
    torch.backends.cudnn.enabled = False
    net = net.cuda()
if num_gpus > 1:
    # 5) 封装
    net = torch.nn.parallel.DistributedDataParallel(net,
                                                   device_ids=[local_rank],
                                                   output_device=local_rank,
                                                find_unused_parameters=True)

Fügen Sie diesen Satz zu den anfänglichen Optionen hinzu

# for distribution
parser.add_argument("--local_rank", type=int)

Läuft beispielsweise auf einem Server mit vier Karten

python -m torch.distributed.launch --nproc_per_node 4 train_ddp.py

Detail

Kontrollpunkt

Speichern Sie einfach das Gewicht des Modells im Hauptthread. Wenn Sie den Zug fortsetzen müssen, müssen Sie das gesamte Prüfpunktwörterbuch von Torch.save speichern. Hier speichere ich vorerst nur das Gewicht des Modells.
Besorgte Freunde können sich auf das Lebenslauftraining des Pytorch-Modells beziehen und das Training auf der Grundlage des Ladens des Modells fortsetzen

if cur_epoch > 0 and cur_epoch % 2 == 0:
    if check_print_rank(opt):
        net_state_dict = net.module.state_dict() if opt.num_gpus > 1 else net.state_dict()
        train_state = {
    
    
            "net": net_state_dict,
            'optimizer': optimizer.state_dict(),
            "cur_epoch": cur_epoch, 
            "cur_step": cur_step
        }
        # 确保路径没错
        train_ckpt_path = 'train_ckpt/save_net_ckpt'
        if not os.path.exists(train_ckpt_path):
            os.makedirs(train_ckpt_path)
        train_state_path = 'train_ckpt/save_state'
        if not os.path.exists(train_state_path):
            os.makedirs(train_state_path)
        
        # 保存模型权重
        torch.save(net_state_dict,
                '{}/epoch_{}.pth'.format(train_ckpt_path, cur_epoch))
        # 保存训练时的状态，方便后续resume train
        torch.save(train_state,
                '{}/state_epoch_{}.pth'.format(train_state_path, cur_epoch))

Follow-up am 14. Juli: Ich konnte mich nicht mehr zurückhalten, der Server hatte einige Probleme und ist direkt abgestorben. Ich muss immer noch ehrlich einen Lebenslauf schreiben ...
Fügen Sie zunächst zwei Optionen hinzu, um die Optionen zu analysieren

# for resume train
parser.add_argument("--resume_train", type=bool, default=False)
parser.add_argument("--resume_state_path", type=str, default='')

Laden Sie dann einige Modellgewichte, Optimierergewichte sowie die aktuelle Epoche und den aktuellen Schritt, bevor das Training beginnt

if opt.resume_train:
    if opt.resume_state_path is not None and opt.resume_state_path != '':
         # 加载之前保存的断点
          state_dict = torch.load(opt.resume_state_path)
          print('load state dict from {}'.format(opt.resume_state_path))
          net.module.load_state_dict(state_dict['net'])
          optimizer.load_state_dict(state_dict['optimizer'])
          cur_epoch = state_dict['cur_epoch']
          cur_step = state_dict['cur_step']
          # -----------------------
          if check_print_rank(opt):
              print('start resume train at epoch {}, step {}'.format(cur_epoch, cur_step))
else:
    optimizer = optim.AdamW(net.parameters(), lr=2e-4)
    cur_epoch = 0
    cur_step = 0
    if check_print_rank(opt):
        print('train from epoch 0, step 0 ...')

Stellen Sie dann den Planer für die Lernratenoptimierung ein

scheduler = LinearWarmupCosineAnnealingLR(
optimizer=optimizer, warmup_epochs=15, max_epochs=150)
max_epoch = opt.epochs

while cur_epoch <= max_epoch:
     cur_epoch = cur_epoch + 1
     scheduler.step(cur_epoch)
     if check_print_rank(opt):
            # lr = scheduler.get_lr()
            lr = scheduler.get_last_lr()[0]
            print('--> cur_epoch: {}, use lr: {}\n'.format(cur_epoch, lr))
            writer.add_scalar("learning rate", lr, cur_epoch)

     for batch in tqdm(trainloader):
         net.train()
         cur_step = cur_step + 1
         optimizer.zero_grad()
         # ... 训练的代码

Speichern Sie Modellgewichte und -status

if (cur_epoch + 1) % 2 == 0:
	net_state_dict = net.module.state_dict() if opt.num_gpus > 1 else net.state_dict()
	  train_state = {
    
    
	      "net": net_state_dict,
	      'optimizer': optimizer.state_dict(),
	      "cur_epoch": cur_epoch, 
	      "cur_step": cur_step
	  }
	  
	if check_print_rank(opt):
	    # 确保路径没错
	    train_ckpt_path = 'train_ckpt/save_net_ckpt'
	    if not os.path.exists(train_ckpt_path):
	        os.makedirs(train_ckpt_path)
	    train_state_path = 'train_ckpt/save_state'
	    if not os.path.exists(train_state_path):
	        os.makedirs(train_state_path)
	    
	    # 保存模型权重
	    torch.save(net_state_dict,
	            '{}/epoch_{}.pth'.format(train_ckpt_path, cur_epoch))
	    # 保存训练时的状态，方便后续resume train
	    torch.save(train_state,
	            '{}/state_epoch_{}.pth'.format(train_state_path, cur_epoch))

Validierung

Hier verwende ich nur den Hauptthread mit dem Rang 0, um die Überprüfung durchzuführen, da es etwas mühsam ist, die Testergebnisse zu synthetisieren, wenn der Multi-Card-Test der generativen Aufgabe durchgeführt wird (natürlich ist das nicht unmöglich, aber ich Ich bin faul), also werde ich es nicht tun. Wenn die Klassifizierungsaufgabe bequemer ist.

if opt.num_gpus <= 1:
    print('epoch: {}, start eval...'.format(cur_epoch))
    eval_operation(net, writer, cur_epoch)
    print('epoch: {}, eval end\n\n'.format(cur_epoch))
else:
    if torch.distributed.get_rank() == 0:
        print('\n rank: {}, epoch: {}, start eval...'.format(torch.distributed.get_rank(), cur_epoch))
        eval_operation(net.module, writer, cur_epoch)
        print('\n rank: {}, epoch: {}, eval end\n\n'.format(torch.distributed.get_rank(), cur_epoch))

Tensorboard

Tensorboard verwendet häufig
Remote-Anzeigevorgänge. Sie können meinen anderen Blog lesen, in dem Tensorboard Remote-Server anzeigt

def check_print_rank(opt):
    return opt.num_gpus <= 1 or torch.distributed.get_rank() == 0

if check_print_rank(opt):
    # logs存在Tensorboard_logs目录下
    start_time = datetime.datetime.now().strftime('%Y-%m-%d  %H:%M:%S')
    Tensorboard_logs_dir = "Tensorboard_logs/{}".format(start_time)
    writer = SummaryWriter(Tensorboard_logs_dir)
    print('log in {}...\n'.format(Tensorboard_logs_dir))

opt hier ist dieses Ding

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--num_gpus",type=int,default= 4,help = "Number of GPUs to use for training")
# ... 一系列命令行参数配置
opt = parser.parse_args()

Verwenden Sie beim Schreiben einfach den Hauptthread zum Schreiben

if check_print_rank(opt):
      # Logging to TensorBoard (if installed) by default
      writer.add_scalar("train loss", loss, cur_step)

Nach Gebrauch endlich verschließen

if check_print_rank(opt):
	writer.close()