train_distributed.py

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
import torchvision
from torch.utils.data.distributed import DistributedSampler

from sacred import Experiment, cli_option

from warmup_scheduler import GradualWarmupScheduler

from lib.datasets import ds
from lib.datasets import StaticHdF5Dataset
from lib.model import net
from lib.model import IODINE
from lib.geco import GECO
from lib.visualization import visualize_slots

from tqdm import tqdm
from pathlib import Path
import shutil
import pprint
import os

import numpy as np
from collections import deque
import time

@cli_option('-r','--local_rank')
def local_rank_option(args, run):
    run.info['local_rank'] = args

ex = Experiment('TRAINING', ingredients=[ds, net], additional_cli_options=[local_rank_option])

@ex.config
def cfg():
    training = {
            'DDP_port': 29500,
            'batch_size': 16,  # training mini-batch size
            'num_workers': 8,  # pytorch dataloader workers
            'iters': 500000,  # train steps if no curriculum
            'lr': 3e-4,  # Adam LR
            'warmup': 10000,
            'decay_rate': 0.5,
            'decay_steps': 100000,   
            'use_scheduler': True,         
            'use_geco': True,
            'geco_reconstruction_target': -23000,  # GECO C
            'geco_ema_alpha': 0.99,  # GECO EMA step parameter,
            'geco_beta_stepsize': 1e-6,  # GECO Lagrange parameter beta
            'mode': 'train',
            'tensorboard_freq': 100,  # how often to write to TB
            'tensorboard_delete_prev': False,  # delete TB dir if already exists
            'checkpoint_freq': 25000,  # save checkpoints every % steps
            'load_from_checkpoint': False,  # whether to load from a checkpoint or not
            'checkpoint': '',  # name of .pth file to load model state
            'run_suffix': 'debug',  # string to append to run name
            'out_dir': 'experiments'
        }

def save_checkpoint(step, model, model_opt, filepath):
    state = {
        'step': step,
        'model': model.state_dict(),
        'model_opt': model_opt.state_dict(),
    }
    torch.save(state, filepath)

@ex.automain
def run(training, seed, _run):
    # maybe create
    run_dir = Path(training['out_dir'], 'runs')
    checkpoint_dir = Path(training['out_dir'], 'weights')
    tb_dir = Path(training['out_dir'], 'tb')
    
    for dir_ in [run_dir, checkpoint_dir, tb_dir]:
        if not dir_.exists():
            #dir_.mkdir()
            print(f'Create {dir_} before running!')
            exit(1)

    tb_dbg = tb_dir / training['run_suffix']

    local_rank = 'cuda:{}'.format(_run.info['local_rank'])
    if local_rank == 'cuda:0':
        writer = SummaryWriter(tb_dbg)
    
    # Fix random seed
    print(f'setting random seed to {seed}')
    # Auto-set by sacred
    # torch.manual_seed(seed)
    torch.backends.cudnn.deterministic=True
    torch.backends.cudnn.benchmark=False
    # Auto-set by sacred 
    #np.random.seed(seed)
    
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = str(training['DDP_port'])
    torch.distributed.init_process_group(backend='nccl')
    
    torch.cuda.set_device(local_rank)

    model = IODINE(batch_size=training['batch_size'])

    if training['use_geco']:
        model_geco = GECO(training['geco_reconstruction_target'], training['geco_ema_alpha'])
    else:
        model_geco = None

    model = model.to(local_rank)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
    model.train()

    print(f'Model parameters #: {torch.nn.utils.parameters_to_vector(model.parameters()).shape[0]}')

    # Optimization
    model_opt = torch.optim.Adam(model.parameters(), lr=training['lr'])
    if training['use_scheduler']:
        scheduler = torch.optim.lr_scheduler.LambdaLR(model_opt, lr_lambda=lambda epoch: 0.5 ** (epoch / 100000))
        scheduler_warmup = GradualWarmupScheduler(model_opt, multiplier=1, total_epoch=training['warmup'], after_scheduler=scheduler)
    else:
        scheduler_warmup = None

    if not training['load_from_checkpoint']:    
        step = 0 
        checkpoint_step = 0
    else:
        checkpoint = checkpoint_dir / training['checkpoint']
        map_location = {'cuda:0': local_rank}
        state = torch.load(checkpoint, map_location=map_location)
        model.load_state_dict(state['model'])
        model_opt.load_state_dict(state['model_opt'])
        step = state['step']
        checkpoint_step = step

    tr_dataset = StaticHdF5Dataset(d_set=training['mode'])
    batch_size = training['batch_size']
    tr_sampler = DistributedSampler(dataset=tr_dataset)
    tr_dataloader = torch.utils.data.DataLoader(tr_dataset,
            batch_size=batch_size, sampler=tr_sampler, 
            num_workers=training['num_workers'], drop_last=True)
    
    max_iters = training['iters']
    
    forward_queue = deque(maxlen=10000)
    backward_queue = deque(maxlen=10000)

    while step <= max_iters:
        
        if local_rank == 'cuda:0':
            data_iter = tqdm(tr_dataloader)
        else:
            data_iter = tr_dataloader

        for batch in data_iter:
            img_batch = batch['imgs'].to(local_rank)
            model_opt.zero_grad()
            
            start = time.time()
            out_dict = model(img_batch, model_geco, step)
            forward_queue.append(time.time() - start)

            total_loss = out_dict['total_loss']
            kl = out_dict['kl']
            nll = out_dict['nll']
            
            start = time.time()
            total_loss.backward()
            backward_queue.append(time.time() - start)

            if training['use_scheduler']:
                scheduler_warmup.step(step)

            # clip gradient norm to 5
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.)

            model_opt.step()

            if training['use_geco'] and model.module.kl_beta > 0.:
                if step == model.module.geco_warm_start:
                    model.module.geco_C_ema = model_geco.init_ema(model.module.geco_C_ema, nll)
                elif step > model.module.geco_warm_start:
                    model.module.geco_C_ema = model_geco.update_ema(model.module.geco_C_ema, nll)
                    model.module.geco_beta = model_geco.step_beta(model.module.geco_C_ema,
                            model.module.geco_beta, training['geco_beta_stepsize'])

            # logging
            if step % training['tensorboard_freq'] == 0 and local_rank == 'cuda:0':
                writer.add_scalar('train/total_loss', total_loss, step)
                writer.add_scalar('train/KL', kl, step)
                writer.add_scalar('train/NLL', nll, step)
                visualize_slots(writer, img_batch, out_dict, step)
                if training['use_geco']:
                    writer.add_scalar('train/geco_beta', model.module.geco_beta, step)
                    writer.add_scalar('train/geco_C_ema', model.module.geco_C_ema, step)

                print('forward time (ms): {}'.format(np.mean(forward_queue) * 1000.))
                print('backward time (ms): {}'.format(np.mean(backward_queue) * 1000.))

            if step > 0 and step % training['checkpoint_freq'] == 0 and local_rank == 'cuda:0':
                prefix = training['run_suffix']
                save_checkpoint(step, model, model_opt, 
                       checkpoint_dir / f'{prefix}-state-{step}.pth')
            
            if step >= max_iters:
                step += 1
                break
            step += 1

    if local_rank == 'cuda:0':
        writer.close()