# --------------------------------------------------------
# X-Decoder -- Generalized Decoding for Pixel, Image, and Language
# Copyright (c) 2022 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Modified by Xueyan Zou (xueyan@cs.wisc.edu)
# --------------------------------------------------------
from datetime import datetime
import time
import os
import sys
import importlib
import json
import random
#import wandb
import logging
import numpy as np
import copy
import contextlib
import shutil
from typing import Any, Callable, Union
import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from mpi4py import MPI
from infinibatch import iterators
from .distributed_trainer import DistributedTrainer
from .utils_trainer import UtilsTrainer
from .utils.misc import *
from .utils.serialization import JSONEncoder, filter_jsonable
logger = logging.getLogger(__name__)
class DefaultTrainer(UtilsTrainer, DistributedTrainer):
def __init__(self, opt):
"""
Set up the task the model is being trained for.
"""
super().__init__(opt)
base_name = 'base_dir'
base_path = os.path.join(self.opt['base_path'], '__init__.py')
spec = importlib.util.spec_from_file_location(base_name, base_path)
module = importlib.util.module_from_spec(spec)
sys.modules[base_name] = module
spec.loader.exec_module(module)
logger.info(f"Imported {base_name} at base_path {self.opt['base_path']}")
pipeline_module = importlib.import_module(f"base_dir.pipeline.{self.opt['PIPELINE']}")
pipeline_class = getattr(pipeline_module, self.opt['PIPELINE'])
logger.info(f"Pipeline for training: {self.opt['PIPELINE']}")
self.pipeline = pipeline_class(self.opt)
def eval(self, ):
logger.info('-----------------------------------------------')
logger.info("Evaluating model ... ")
self.mode = "eval"
# self.model_names, self.raw_models, self.criteria = self.pipeline.set_up_model()
self.raw_models = self.pipeline.initialize_model()
self.model_names = self.raw_models.keys()
# move models to the device
for module_name in self.model_names:
self.raw_models[module_name].to(self.opt['device'])
# load model during evaluation
if self.opt['WEIGHT'] and os.path.isfile(self.opt['RESUME_FROM']):
model_path = self.opt['RESUME_FROM']
self.load_model(model_path)
else:
raise ValueError(f"Model not found: {model_path}")
results = self._eval_on_set(self.save_folder)
return results
def _eval_on_set(self, save_folder):
logger.info(f"Evaluation start ...")
if self.opt['FP16']:
from torch.cuda.amp import autocast
with autocast():
results = self.pipeline.evaluate_model(self, save_folder)
else:
results = self.pipeline.evaluate_model(self, save_folder)
if self.opt['rank'] == 0:
logger.info(results)
return results
def compute_loss(self, forward_func, batch):
def forward(func, trainer, batch):
if self.opt['FP16']:
from torch.cuda.amp import autocast
with autocast():
loss = func(trainer, batch)
else:
loss = func(trainer, batch)
return loss
loss = forward(forward_func, self, batch)
return loss
def backward_loss(self, loss, model_names=['default']): # noqa: E252
def backward(loss_tensor):
if self.opt['FP16']:
self.grad_scaler.scale(loss_tensor).backward()
else:
loss_tensor.backward()
if self.grad_acc_steps > 1:
loss = loss / self.grad_acc_steps
backward(loss)
return loss
def update_model(self, model_name='default'):
if self.opt['FP16']:
self.grad_scaler.unscale_(self.optimizers[model_name])
self.grad_scaler.step(self.optimizers[model_name])
else:
self.optimizers[model_name].step()
self.optimizers[model_name].zero_grad()
self.train_params['optim_steps'][model_name] += 1
self.lr_schedulers[model_name].step()
def train_step(self, batch):
self.grad_acc_batches.append(batch) # support batch accumulation
if self.is_gradient_accumulation_boundary():
# set all modules and criteria into training mode
for model_name in self.model_names:
self.models[model_name].train()
assert len(self.grad_acc_batches) == self.grad_acc_steps
total_batch_sample = 0
for batch_index, batch in enumerate(self.grad_acc_batches):
loss_info, sample_size_info, extra_info = \
self.pipeline.forward_step(self,
batch,
self.grad_acc_batches,
batch_index,
is_distributed=(self.opt['world_size'] > 1))
self.train_loss.update_iter(loss_info)
total_batch_sample += sample_size_info['num_samples']
if self.opt['FP16']:
# Update GradScaler after an effective batch
self.grad_scaler.update()
# update losses and item counts of an effective batch to the AverageMeters
if self.opt['world_size'] > 1:
total_batch_sample = torch.tensor(total_batch_sample).to(self.opt['device'])
torch.distributed.all_reduce(total_batch_sample, torch.distributed.ReduceOp.SUM)
total_batch_sample = total_batch_sample.item()
self.train_params['total_batch_size'] += total_batch_sample
self.grad_acc_batches = []
self.train_params['num_updates'] += 1
def init_train(self):
self.mode = "train"
logger.info('-------------------------------------------------------')
logger.info("Training on rank: {}".format(self.opt['rank']))
self.raw_models = self.pipeline.initialize_model()
self.model_names = list(self.raw_models.keys())
# move models to the device
for module_name in self.model_names:
self.raw_models[module_name].to(self.opt['device'])
self.train_dataloaders = self.pipeline.get_dataloaders(self, 'train', is_evaluation=False)
self.train_params = {
"updates_per_epoch": len(self.train_dataloaders),
"total_batch_size": 0,
"num_updates": 0,
"optim_steps": {module_name: 0 for module_name in self.model_names},
"start_epoch_idx": 0,
"start_batch_idx": 0,
"current_epoch_idx": 0,
"current_batch_idx": 0,
"resume_epoch_idx": 0,
}
self.train_loss = LossMeter()
self.grad_acc_batches = []
if self.opt['CUDA']:
torch.cuda.empty_cache()
self.create_optimizer_and_scheduler()
self.models = {model_name: self.raw_models[model_name] for model_name in self.model_names}
self._initialize_ddp()
if self.opt.get('WEIGHT', False):
self.load_weight(self.opt['RESUME_FROM'], must_exist=True)
if self.opt.get('RESUME', False):
self.load_checkpoint(self.opt['RESUME_FROM'], must_exist=True)
######################
# Start the main loop
######################
if self.opt['rank'] == 0:
# Train!
logger.info("***** Running training *****")
logger.info(f" Num of GPUs = {self.opt['world_size']}")
logger.info(f" Num Epochs = {self.opt['SOLVER']['MAX_NUM_EPOCHS']}")
logger.info(f" Num of Mini Batches per Epoch = {self.train_params['updates_per_epoch']}")
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {self.opt['SOLVER']['MAX_NUM_EPOCHS'] * self.train_params['updates_per_epoch']}")
logger.info(f" Gradient Accumulation steps = {self.grad_acc_steps}")
logger.info(f" Total optimization steps = {self.opt['SOLVER']['MAX_NUM_EPOCHS'] * self.train_params['updates_per_epoch'] // self.grad_acc_steps}")
def train(self):
"""
Training
"""
self.init_train()
current_optim_steps = self._get_and_validate_current_optim_steps()
num_epochs = self.opt['SOLVER']['MAX_NUM_EPOCHS']
if self.opt.get('EVAL_AT_START', False):
results = self._eval_on_set(self.save_folder)
# if self.opt['rank'] == 0 and self.opt['WANDB']:
# wandb.log(results)
train_prev_logged_time = datetime.now()
for epoch in range(self.train_params['start_epoch_idx'], num_epochs):
self.train_params['current_epoch_idx'] = epoch
logger.info(f"Start epoch: {epoch} training.")
epoch_start_time = datetime.now()
for batch_idx, batch in enumerate(self.train_dataloaders):
if self.train_params['current_epoch_idx'] == self.train_params['start_epoch_idx']:
if batch_idx < self.train_params['start_batch_idx']: # skip the first few batches for resuming
continue
self.train_params['current_batch_idx'] = batch_idx
prev_optim_steps = current_optim_steps
prev_total_batch_size = self.train_params['total_batch_size']
# update
self.prev_optim_steps = prev_optim_steps
self.train_step(batch)
current_optim_steps = self._get_and_validate_current_optim_steps()
# logging
if prev_optim_steps != current_optim_steps: # an optimizer update was made
log_first = self.opt.get("LOG_FIRST", 10)
log_every = self.opt.get("LOG_EVERY", 100)
if (current_optim_steps % log_every == 0) or (epoch == 0 and current_optim_steps <= log_first): # print logging
last_lr = {}
for module_name in self.model_names:
last_lr[module_name] = self.lr_schedulers[module_name].get_last_lr()[0]
train_time_delta = (datetime.now() - train_prev_logged_time).total_seconds()
train_prev_logged_time = datetime.now()
MB = 1024.0 * 1024.0
memory = torch.cuda.max_memory_allocated() / MB
if self.opt['rank'] == 0:
# if self.opt['WANDB']:
# # log for wandb
# wb_loss_info = {key: obj.val for key, obj in self.train_loss.losses.items()}
# wandb.log(wb_loss_info, step=self.prev_optim_steps)
# log for terminal
logger.info(f"epochs[{epoch:6}] optim steps[{current_optim_steps:.0f}] "
f"learning rate[{', '.join([f'{key}: {val:.5e}' for key, val in last_lr.items()])}] "
f"train loss[{', '.join([f'{key}: {obj.val:.5f}/{obj.avg:.5f}' for key, obj in self.train_loss.losses.items()])}] "
# f"total_loss[{total_loss:.5f}/{total_loss_avg:.5f} "
f"items per batch[{self.train_params['total_batch_size'] - prev_total_batch_size}] "
f"items per second[{(self.train_params['total_batch_size'] - prev_total_batch_size) / train_time_delta:.2f}] "
f"total items[{self.train_params['total_batch_size']}] "
f"mini batches[{self.train_params['num_updates']:6}] "
f"memory[{memory:.0f}] "
f"epoch remaining[{str((datetime.now() - epoch_start_time) / (batch_idx + 1) * (self.train_params['updates_per_epoch'] - batch_idx - 1)).split('.')[0]}]")
# evaluate and save ckpt every epoch
if batch_idx + 1 == self.train_params['updates_per_epoch']:
if self.opt.get('SAVE_CHECKPOINT', True):
self.save_checkpoint(self.train_params['num_updates'])
results = self._eval_on_set(self.save_folder)
# if self.opt['rank'] == 0 and self.opt['WANDB']:
# wandb.log(results)
break
logger.info(f"This epoch takes {datetime.now() - epoch_start_time}")
logger.info(f"PROGRESS: {100.0 * (epoch + 1) / num_epochs:.2f}%")
logger.info(f"Config files are at {self.opt['conf_files']}")
# if not self.opt.get('SAVE_CHECKPOINT', True):
# self.save_checkpoint(self.train_params['num_updates'])