280 lines
13 KiB
Python
280 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
|
# @Author: Weisen Pan
|
|
|
|
import torch
|
|
import argparse
|
|
import warnings
|
|
import setproctitle
|
|
from torch import nn, decentralized # Used for decentralized training
|
|
from torch.backends import cudnn # Optimizes performance for convolutional networks
|
|
from tensorboardX import SummaryWriter # For logging metrics and results to TensorBoard
|
|
import torch.cuda.amp as amp # For mixed precision training
|
|
from config import * # Custom configuration module
|
|
from params import train_params # Training parameters
|
|
from utils import label_smoothing, norm, summary, metric, lr_scheduler, prefetch # Utility functions
|
|
from model import coremodel # Core model implementation
|
|
from dataset import factory # Dataset and data loader factory
|
|
from params.train_params import save_hp_to_json # Function to save hyperparameters to JSON
|
|
|
|
# Global variable to store the best accuracy obtained during training
|
|
best_acc1 = 0
|
|
|
|
def main(args):
|
|
# Warn if a specific GPU is chosen, as this will disable data parallelism
|
|
if args.gpu is not None:
|
|
warnings.warn("Selecting a specific GPU will disable data parallelism.")
|
|
|
|
# Adjust loop factor based on specific training configurations
|
|
args.loop_factor = 1 if args.is_train_sep or args.is_single_branch else args.split_factor
|
|
# Check if decentralized training is needed
|
|
args.is_decentralized = args.world_size > 1 or args.multiprocessing_decentralized
|
|
|
|
# Get the number of available GPUs on the machine
|
|
num_gpus = torch.cuda.device_count()
|
|
args.ngpus_per_node = num_gpus
|
|
print(f"INFO:PyTorch: GPUs available on this node: {num_gpus}")
|
|
|
|
# If multiprocessing is needed for decentralized training
|
|
if args.multiprocessing_decentralized:
|
|
# Adjust world size to account for multiple GPUs
|
|
args.world_size *= num_gpus
|
|
# Spawn multiple processes for each GPU
|
|
torch.multiprocessing.spawn(execute_worker_process, nprocs=num_gpus, args=(num_gpus, args))
|
|
else:
|
|
# If using a single GPU
|
|
print("INFO:PyTorch: Using GPU 0 for single GPU training")
|
|
args.gpu = 0
|
|
# Call main worker for single GPU
|
|
execute_worker_process(args.gpu, num_gpus, args)
|
|
|
|
def execute_worker_process(gpu, num_gpus, args):
|
|
global best_acc1
|
|
args.gpu = gpu
|
|
# Set the directory where models will be saved
|
|
args.model_dir = os.path.join(HOME, "models", "coremodel", str(args.spid))
|
|
|
|
# Initialize the decentralized training process group if needed
|
|
if args.is_decentralized:
|
|
print("INFO:PyTorch: Initializing process group for decentralized training.")
|
|
if args.dist_url == "env://" and args.rank == -1:
|
|
args.rank = int(os.environ["RANK"])
|
|
if args.multiprocessing_decentralized:
|
|
args.rank = args.rank * num_gpus + gpu
|
|
decentralized.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank)
|
|
|
|
# Set the GPU to be used for training or evaluation
|
|
if args.gpu is not None:
|
|
print(f"INFO:PyTorch: GPU {args.gpu} in use for training (Rank: {args.rank})" if not args.evaluate else f"INFO:PyTorch: GPU {args.gpu} in use for evaluation (Rank: {args.rank})")
|
|
|
|
# Set process title for better identification in system process monitors
|
|
setproctitle.setproctitle(f"{args.proc_name}centralized_rank{args.rank}")
|
|
|
|
# Initialize a SummaryWriter for TensorBoard logging
|
|
val_writer = SummaryWriter(log_dir=os.path.join(args.model_dir, 'val'))
|
|
|
|
# Use label smoothing if enabled, otherwise use standard cross-entropy loss
|
|
criterion = label_smoothing.label_smoothing_CE(reduction='mean') if args.is_label_smoothing else nn.CrossEntropyLoss()
|
|
|
|
# Instantiate the model
|
|
model = coremodel.coremodel(args, norm_layer=norm.norm(args.norm_mode), criterion=criterion)
|
|
print(f"INFO:PyTorch: Model '{args.arch}' has {metric.get_the_number_of_params(model)} parameters")
|
|
|
|
# If summary is requested, print model and exit
|
|
if args.is_summary:
|
|
print(model)
|
|
return
|
|
|
|
# Save model configuration and hyperparameters
|
|
summary.save_model_to_json(args, model)
|
|
|
|
# Convert BatchNorm layers to synchronized BatchNorm for decentralized training
|
|
if args.is_decentralized and args.world_size > 1 and args.is_syncbn:
|
|
print("INFO:PyTorch: Converting BatchNorm to SyncBatchNorm")
|
|
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
|
|
|
|
# Set up the model for GPU-based training
|
|
if args.gpu is not None:
|
|
torch.cuda.set_device(args.gpu)
|
|
model.cuda(args.gpu)
|
|
args.batch_size = int(args.batch_size / num_gpus) # Adjust batch size for multiple GPUs
|
|
args.workers = int((args.workers + num_gpus - 1) / num_gpus) # Adjust number of workers
|
|
# Use decentralized data parallel model
|
|
model = nn.parallel.decentralizedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
|
|
else:
|
|
# Use standard DataParallel for multi-GPU training
|
|
model = nn.DataParallel(model).cuda()
|
|
|
|
# Create the optimizer
|
|
optimizer = create_optimizer(args, model)
|
|
# Set up the gradient scaler for mixed precision training, if enabled
|
|
scaler = amp.GradScaler() if args.is_amp else None
|
|
|
|
# If resuming from a checkpoint, load model and optimizer state
|
|
if args.resume:
|
|
load_checkpoint(args, model, optimizer, scaler)
|
|
|
|
cudnn.performance_test = True # Enable cuDNN performance optimizations
|
|
|
|
# Set up data loader parameters
|
|
data_loader_params = {
|
|
'split_factor': args.loop_factor if args.is_diff_data_train else 1,
|
|
'batch_size': args.batch_size,
|
|
'crop_size': args.crop_size,
|
|
'dataset': args.dataset,
|
|
'is_decentralized': args.is_decentralized,
|
|
'num_workers': args.workers,
|
|
'randaa': args.randaa,
|
|
'is_autoaugment': args.is_autoaugment,
|
|
'is_cutout': args.is_cutout,
|
|
'erase_p': args.erase_p,
|
|
}
|
|
|
|
# Get the training and validation data loaders
|
|
train_loader, train_sampler = factory.obtain_data_loader(args.data, split="train", **data_loader_params)
|
|
val_loader = factory.obtain_data_loader(args.data, split="val", batch_size=args.eval_batch_size, crop_size=args.crop_size, num_workers=args.workers)
|
|
|
|
# Set up the learning rate scheduler
|
|
scheduler = lr_scheduler.create_scheduler(args, len(train_loader))
|
|
|
|
# If evaluating, run the validation function and exit
|
|
if args.evaluate:
|
|
validate(val_loader, model, args)
|
|
return
|
|
|
|
# Begin training and evaluation
|
|
train_and_evaluate(train_loader, val_loader, train_sampler, model, optimizer, scheduler, scaler, val_writer, args, num_gpus)
|
|
|
|
# Function to create the optimizer
|
|
def create_optimizer(args, model):
|
|
param_groups = model.parameters() if args.is_wd_all else lr_scheduler.get_parameter_groups(model)
|
|
# Select the optimizer based on input arguments
|
|
if args.optimizer == 'SGD':
|
|
return torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.is_nesterov)
|
|
elif args.optimizer == 'AdamW':
|
|
return torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.999), eps=1e-4, weight_decay=args.weight_decay)
|
|
elif args.optimizer == 'RMSprop':
|
|
return torch.optim.RMSprop(param_groups, lr=args.lr, alpha=0.9, momentum=0.9, weight_decay=args.weight_decay)
|
|
else:
|
|
# Raise error if unsupported optimizer is selected
|
|
raise NotImplementedError(f"Optimizer {args.optimizer} not implemented")
|
|
|
|
# Function to load a checkpoint and resume training
|
|
def load_checkpoint(args, model, optimizer, scaler):
|
|
if os.path.isfile(args.resume):
|
|
print(f"INFO:PyTorch: Loading checkpoint from '{args.resume}'")
|
|
loc = f'cuda:{args.gpu}' if args.gpu is not None else None
|
|
checkpoint = torch.load(args.resume, map_location=loc)
|
|
args.start_epoch = checkpoint['epoch']
|
|
global best_acc1
|
|
best_acc1 = checkpoint['best_acc1']
|
|
model.load_state_dict(checkpoint['state_dict'])
|
|
optimizer.load_state_dict(checkpoint['optimizer'])
|
|
if "scaler" in checkpoint:
|
|
scaler.load_state_dict(checkpoint['scaler'])
|
|
print(f"INFO:PyTorch: Checkpoint loaded, epoch {checkpoint['epoch']}")
|
|
else:
|
|
print(f"INFO:PyTorch: No checkpoint found at '{args.resume}'")
|
|
|
|
# Function to train and evaluate the model over multiple epochs
|
|
def train_and_evaluate(train_loader, val_loader, train_sampler, model, optimizer, scheduler, scaler, val_writer, args, num_gpus):
|
|
for epoch in range(args.start_epoch, args.epochs + 1):
|
|
if args.is_decentralized:
|
|
train_sampler.set_epoch(epoch)
|
|
|
|
train_one_epoch(train_loader, model, optimizer, scheduler, epoch, scaler, val_writer, args)
|
|
|
|
# Evaluate the model every 'eval_per_epoch' epochs
|
|
if (epoch + 1) % args.eval_per_epoch == 0:
|
|
acc_all = validate(val_loader, model, args)
|
|
global best_acc1
|
|
is_best = acc_all[0] > best_acc1 # Track the best accuracy
|
|
best_acc1 = max(acc_all[0], best_acc1)
|
|
# Save the model checkpoint
|
|
save_checkpoint(model, optimizer, scaler, epoch, best_acc1, args, is_best)
|
|
|
|
# Function to perform one training epoch
|
|
def train_one_epoch(train_loader, model, optimizer, scheduler, epoch, scaler, val_writer, args):
|
|
metric_storage = create_metric_storage(args.loop_factor)
|
|
model.train() # Set the model to training mode
|
|
data_loader = prefetch.data_prefetcher(train_loader) # Use data prefetching to improve efficiency
|
|
images, target = data_loader.next()
|
|
|
|
optimizer.zero_grad() # Reset gradients
|
|
while images is not None:
|
|
# Adjust the learning rate based on the scheduler
|
|
scheduler(optimizer, epoch)
|
|
|
|
# Perform forward pass with mixed precision if enabled
|
|
if args.is_amp:
|
|
with amp.autocast():
|
|
ensemble_output, outputs, ce_loss, cot_loss = model(images, target=target, mode='train', epoch=epoch)
|
|
else:
|
|
ensemble_output, outputs, ce_loss, cot_loss = model(images, target=target, mode='train', epoch=epoch)
|
|
|
|
# Calculate total loss and normalize
|
|
total_loss = (ce_loss + cot_loss) / args.iters_to_accumulate
|
|
val_writer.add_scalar('average_training_loss', total_loss, global_step=epoch)
|
|
|
|
# Perform backward pass and update gradients with mixed precision if enabled
|
|
if args.is_amp:
|
|
scaler.scale(total_loss).backward()
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
else:
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
|
|
images, target = data_loader.next() # Fetch the next batch of data
|
|
|
|
# Function to save the model checkpoint
|
|
def save_checkpoint(model, optimizer, scaler, epoch, best_acc1, args, is_best):
|
|
ckpt = {
|
|
'epoch': epoch + 1,
|
|
'state_dict': model.state_dict(),
|
|
'best_acc1': best_acc1,
|
|
'optimizer': optimizer.state_dict(),
|
|
}
|
|
if args.is_amp:
|
|
ckpt['scaler'] = scaler.state_dict()
|
|
metric.save_checkpoint(ckpt, is_best, args.model_dir, filename=f"checkpoint_{epoch}.pth.tar")
|
|
|
|
# Function to validate the model on the validation dataset
|
|
def validate(val_loader, model, args):
|
|
metric_storage = create_metric_storage(args.loop_factor)
|
|
model.eval() # Set the model to evaluation mode
|
|
|
|
with torch.no_grad():
|
|
for i, (images, target) in enumerate(val_loader):
|
|
if args.gpu is not None:
|
|
images = images.cuda(args.gpu, non_blocking=True)
|
|
target = target.cuda(args.gpu, non_blocking=True)
|
|
|
|
# Perform forward pass with mixed precision if enabled
|
|
if args.is_amp:
|
|
with amp.autocast():
|
|
ensemble_output, outputs, ce_loss = model(images, target=target, mode='val')
|
|
else:
|
|
ensemble_output, outputs, ce_loss = model(images, target=target, mode='val')
|
|
|
|
batch_size = images.size(0)
|
|
acc1, acc5 = metric.accuracy(ensemble_output, target, topk=(1, 5))
|
|
|
|
metric_storage.update(acc1, acc5, ce_loss, batch_size)
|
|
|
|
return metric_storage.results()
|
|
|
|
# Helper function to create a storage for metrics during training and validation
|
|
def create_metric_storage(loop_factor):
|
|
# Initialize metrics for accuracy and other performance metrics
|
|
top1_all = [metric.AverageMeter(f'Acc@1_{i}', ':6.2f') for i in range(loop_factor)]
|
|
avg_top1 = metric.AverageMeter('Avg_Acc@1', ':6.2f')
|
|
return metric.ProgressMeter(len(top1_all), top1_all, avg_top1)
|
|
|
|
# Main entry point for the script
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='Centralized Training')
|
|
args = train_params.add_parser_params(parser) # Add parameters to the argument parser
|
|
assert args.is_fed == 0, "Centralized training requires args.is_fed to be False"
|
|
os.makedirs(args.model_dir, exist_ok=True) # Create model directory if it doesn't exist
|
|
main(args) # Call the main function
|