train_ddp.py

# 训练
# 分布式训练+自动混合精度
# CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node 2 --master_port 6036 train_ddp.py --model AdaAtt --id test --batch_size 100 --nproc_per_node 2

import torch
import random
import numpy as np

import time
import torch.nn as nn

from config import config
from data_load import data_load_ddp

from utils.import_models import construct_model
from utils.loss import Cross_Entropy
from utils.log import Log_Writer, train_print
from utils.eval import generate_captions, eval_pycoco

# ddp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

local_rank = config.local_rank
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")
device = torch.device('cuda', local_rank)

# 随机种子
seed = config.seed
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True

# log
if local_rank == 0:
   writer = Log_Writer(config)
global_step = 0
loss_ce_avg = 0

# data_loader
train_loader = data_load_ddp(config, config.train, 'train')

# model
model = construct_model(config).to(device)
if config.step != 0:
    log_path = config.log_dir.format(config.id)
    trained_model_path = log_path + '/model/model_' + str(config.step) + '.pt'
    model.load_state_dict(torch.load(trained_model_path))
    global_step = config.step
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

# optimizer
if config.model in ['NIC', 'SAT', 'AdaAtt']:
    optimizer = torch.optim.Adam([{'params': model.module.resnet_encoder.parameters(), 'lr': config.lr_enc},
                                  {'params': model.module.lstm_decoder_att.parameters(), 'lr': config.learning_rate}],
                                  betas=(0.9, 0.98), eps=1e-9)
else:
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate, betas=(0.9, 0.98), eps=1e-9)

# mixed precision
scaler = torch.cuda.amp.GradScaler()

# loss
loss_fn = Cross_Entropy()

for epoch in range(config.epochs):
    train_loader.sampler.set_epoch(epoch)
    model.train()
    totel_step = len(train_loader)
    epoch_time = time.time()
    step_time = time.time()

    if config.model in ['NIC', 'SAT', 'AdaAtt']:
        if epoch == config.ft_epoch:
            model.module.fine_tune()

    for step, (image_feature, cap, cap_len) in enumerate(train_loader):

        global_step += 1
        optimizer.zero_grad()

        image_feature = {k: v.to(device) for k, v in image_feature.items()}
        cap = cap.to(device)
        cap_len = cap_len.to(device)

        with torch.cuda.amp.autocast():

            logit = model(image_feature, cap, cap_len)

            loss_ce = loss_fn(logit, cap, cap_len)

        loss_ce_avg += loss_ce.item()

        scaler.scale(loss_ce).backward()

        scaler.unscale_(optimizer)
        nn.utils.clip_grad_value_(model.parameters(), config.grad_clip)

        scaler.step(optimizer)

        scaler.update()

        if global_step % config.save_loss_freq == 0:
            if local_rank == 0:
               writer.write_tensorboard('loss_ce', loss_ce_avg/config.save_loss_freq, global_step)
            loss_ce_avg = 0

        train_print(loss_ce.item(), step, totel_step, epoch, time.time() - step_time, time.time() - epoch_time)
        step_time = time.time()

        if local_rank == 0:
            if global_step % config.save_model_freq == 0:
                print("Evaluating...")

                # 保存模型
                writer.save_model(model.module, global_step)

                # validation
                model.module.eval()
                gen_pycoco_path = generate_captions(config, model.module, global_step, 'val')
                pycoco_results = eval_pycoco(config, gen_pycoco_path, 'val')
                writer.write_metrics(pycoco_results, global_step)

                model.module.train()