# -*- coding: utf-8 -*- from __future__ import absolute_import import os import sys import time if os.getcwd() not in sys.path: sys.path.append(os.getcwd()) import argparse import random from datetime import datetime import numpy as np import torch import torch.nn as nn import torch.optim as optim from configs.get_config import load_config from datasets import * from lib.core_function import test, train, validate from lib.optimizers.sam import SAM from lib.scheduler.linear_decay import LinearDecayLR from logs.logger import LOG_DIR, Logger from losses import * from models import * from package_utils.misc import NativeScalerWithGradNormCount as NativeScaler from tensorboardX import SummaryWriter from torch.utils.data import DataLoader def args_parser(args=None): parser = argparse.ArgumentParser("Training process...") parser.add_argument("--cfg", help="Config file", required=True) parser.add_argument( "--alloc_mem", "-a", help="Pre allocating GPU memory", action="store_true" ) return parser.parse_args(args) if __name__ == "__main__": if len(sys.argv[1:]): args = sys.argv[1:] else: args = None args = args_parser(args) cfg = load_config(args.cfg) logger = Logger(task=f"training_{cfg.TASK}") # Seed seed = cfg.SEED random.seed(seed) torch.manual_seed(seed) np.random.seed(seed) torch.cuda.manual_seed(seed) # Allocate memory if args.alloc_mem: mem_all_tensors = torch.rand(60, 10000, 10000) mem_all_tensors.to("cuda:0") # Configuing GPU devices devices = torch.device("cpu") torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True if "gpus" in cfg.TRAIN.gpus and cfg.TRAIN.gpus is not None: # Only support a single gpu for training now devices = torch.device("cuda:1") model = build_model(cfg.MODEL, MODELS).cuda() # Loading Dataloader start_loading = time.time() val_dataset = build_dataset( cfg.DATASET, DATASETS, default_args=dict(split="val", config=cfg.DATASET) ) val_dataloader = DataLoader( val_dataset, batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), shuffle=False, pin_memory=cfg.DATASET.PIN_MEMORY, num_workers=cfg.DATASET.NUM_WORKERS, worker_init_fn=val_dataset.train_worker_init_fn, collate_fn=val_dataset.train_collate_fn, ) logger.info( "Loading val dataloader successfully! -- {}".format(time.time() - start_loading) ) start_loading = time.time() train_dataset = build_dataset( cfg.DATASET, DATASETS, default_args=dict(split="train", config=cfg.DATASET) ) train_dataloader = DataLoader( train_dataset, batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), shuffle=True, pin_memory=cfg.DATASET.PIN_MEMORY, num_workers=cfg.DATASET.NUM_WORKERS, worker_init_fn=train_dataset.train_worker_init_fn, collate_fn=train_dataset.train_collate_fn, ) logger.info( "Loading Train dataloader successfully! -- {}".format( time.time() - start_loading ) ) # Defining Loss function and Optimizer critetion = build_losses( cfg.TRAIN.loss, LOSSES, default_args=dict(cfg=cfg.TRAIN.loss) ).cuda() if cfg.TRAIN.use_amp: eff_lr = ( cfg.TRAIN.lr * cfg.TRAIN.accumulation_steps * cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus) / 64 ) # 16*4=64 as default, might change else: eff_lr = cfg.TRAIN.lr if cfg.TRAIN.optimizer == "Adam": optimizer = optim.Adam(model.parameters(), lr=eff_lr, weight_decay=1e-4) elif cfg.TRAIN.optimizer == "AdamW": optimizer = optim.AdamW( model.parameters(), lr=eff_lr, betas=(0.9, 0.999), weight_decay=1e-4 ) elif cfg.TRAIN.optimizer == "SAM": # optimizer = SAM(model.parameters(), optim.Adam, lr=cfg.TRAIN.lr, weight_decay=1e-4) optimizer = SAM( model.parameters(), optim.Adam, lr=eff_lr, betas=(0.9, 0.995), weight_decay=1e-4, ) else: optimizer = optim.SGD( model.parameters(), lr=eff_lr, weight_decay=1e-5, momentum=0.9 ) # Defining scaler scaler = NativeScaler() if cfg.TRAIN.use_amp else None # Loading model model, optimizer, start_epoch, scaler = preset_model( cfg, model, optimizer=optimizer, scaler=scaler ) if len(cfg.TRAIN.gpus) > 0: model = nn.DataParallel(model, device_ids=cfg.TRAIN.gpus).cuda() else: model = model.cuda() # Learning rate Scheduler if cfg.TRAIN.lr_scheduler == "MultiStepLR": lr_scheduler = optim.lr_scheduler.MultiStepLR( optimizer, **cfg.TRAIN.lr_scheduler ) else: lr_scheduler = LinearDecayLR( optimizer, cfg.TRAIN.epochs, cfg.TRAIN.epochs // cfg.TRAIN.start_decay, last_epoch=cfg.TRAIN.begin_epoch, booster=cfg.TRAIN.booster, ) # Enabling tensorboard writer = SummaryWriter( ".tensorboard/{}_{}".format(datetime.today().strftime("%Y-%m-%d"), cfg.TASK) ) trainIters = 0 valIters = 0 min_val_loss = 1e10 max_val_acc = 0 max_test_auc = 0 metrics_base = ( cfg.METRICS_BASE ) # Combine heatmap + cls prediction to calculate accuracy # Starting training process logger.info("Starting training process...") for epoch in range(start_epoch, cfg.TRAIN.epochs): # Unfreezin backbone to update weights if cfg.TRAIN.freeze_backbone and epoch == cfg.TRAIN.warm_up: unfreeze_backbone(model) np.random.seed(seed + epoch) if epoch > 0 and cfg.DATA_RELOAD: logger.info(f"Reloading data for epoch {epoch}...") train_dataset._reload_data(epoch=epoch) train_dataloader = DataLoader( train_dataset, batch_size=cfg.TRAIN.batch_size * len(cfg.TRAIN.gpus), shuffle=True, pin_memory=cfg.DATASET.PIN_MEMORY, num_workers=cfg.DATASET.NUM_WORKERS, worker_init_fn=train_dataset.train_worker_init_fn, collate_fn=train_dataset.train_collate_fn, ) loss_avg, acc_avg, trainIters = train( cfg, model, critetion, optimizer, epoch, train_dataloader, logger, writer, devices, trainIters, metrics_base=metrics_base, scaler=scaler, ) if epoch % cfg.TRAIN.every_val_epochs == 0: loss_val, acc_val, valIters = validate( cfg, model, critetion, epoch, val_dataloader, logger, writer, devices, valIters, metrics_base=metrics_base, ) if acc_val.avg > max_val_acc: # Saving checkpoint ckp_path = os.path.join( LOG_DIR, "{}_{}_model_best.pth".format(cfg.MODEL.type, cfg.TASK) ) save_model(path=ckp_path, epoch=epoch, model=model, optimizer=optimizer) min_val_loss = loss_val.avg max_val_acc = acc_val.avg logger.info(f"Saved best model at epoch --- {epoch}") lr_scheduler.step()