import os import os.path as osp import torch import torch.nn as nn from tqdm import tqdm import torch.distributed as dist import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler from model.loss import FocalLoss from tools.dataset import load_data, MultiEpochsDataLoader import matplotlib.pyplot as plt from configs import trainer_tools import yaml from datetime import datetime def load_configuration(config_path='configs/compare.yml'): """加载配置文件""" with open(config_path, 'r') as f: return yaml.load(f, Loader=yaml.FullLoader) def initialize_model_and_metric(conf, class_num): """初始化模型和度量方法""" tr_tools = trainer_tools(conf) backbone_mapping = tr_tools.get_backbone() metric_mapping = tr_tools.get_metric(class_num) if conf['models']['backbone'] in backbone_mapping: model = backbone_mapping[conf['models']['backbone']]() else: raise ValueError('不支持该模型: {}'.format({conf['models']['backbone']})) if conf['training']['metric'] in metric_mapping: metric = metric_mapping[conf['training']['metric']]() else: raise ValueError('不支持的metric类型: {}'.format(conf['training']['metric'])) return model, metric def setup_optimizer_and_scheduler(conf, model, metric): """设置优化器和学习率调度器""" tr_tools = trainer_tools(conf) optimizer_mapping = tr_tools.get_optimizer(model, metric) if conf['training']['optimizer'] in optimizer_mapping: optimizer = optimizer_mapping[conf['training']['optimizer']]() scheduler_mapping = tr_tools.get_scheduler(optimizer) scheduler = scheduler_mapping[conf['training']['scheduler']]() print('使用{}优化器 使用{}调度器'.format(conf['training']['optimizer'], conf['training']['scheduler'])) return optimizer, scheduler else: raise ValueError('不支持的优化器类型: {}'.format(conf['training']['optimizer'])) def setup_loss_function(conf): """配置损失函数""" if conf['training']['loss'] == 'focal_loss': return FocalLoss(gamma=2) else: return nn.CrossEntropyLoss() def train_one_epoch(model, metric, criterion, optimizer, dataloader, device, scaler, conf): """执行单个训练周期""" model.train() train_loss = 0 for data, labels in tqdm(dataloader, desc="Training", ascii=True, total=len(dataloader)): data = data.to(device) labels = labels.to(device) # with torch.cuda.amp.autocast(): embeddings = model(data) if not conf['training']['metric'] == 'softmax': thetas = metric(embeddings, labels) else: thetas = metric(embeddings) loss = criterion(thetas, labels) optimizer.zero_grad() scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() train_loss += loss.item() return train_loss / len(dataloader) def validate(model, metric, criterion, dataloader, device, conf): """执行验证""" model.eval() val_loss = 0 with torch.no_grad(): for data, labels in tqdm(dataloader, desc="Validating", ascii=True, total=len(dataloader)): data = data.to(device) labels = labels.to(device) embeddings = model(data) if not conf['training']['metric'] == 'softmax': thetas = metric(embeddings, labels) else: thetas = metric(embeddings) loss = criterion(thetas, labels) val_loss += loss.item() return val_loss / len(dataloader) def save_model(model, path, is_parallel): """保存模型权重""" if is_parallel: torch.save(model.module.state_dict(), path) else: torch.save(model.state_dict(), path) def log_training_info(log_path, log_info): """记录训练信息到日志文件""" with open(log_path, 'a') as f: f.write(log_info + '\n') def initialize_training_components(distributed=False): """初始化所有训练所需组件""" # 加载配置 conf = load_configuration() # 初始化分布式训练相关参数 components = { 'conf': conf, 'distributed': distributed, 'device': None, 'train_dataloader': None, 'val_dataloader': None, 'model': None, 'metric': None, 'criterion': None, 'optimizer': None, 'scheduler': None, 'checkpoints': None, 'scaler': None } # 如果是非分布式训练，直接创建所有组件 if not distributed: # 数据加载 train_dataloader, class_num = load_data(training=True, cfg=conf, return_dataset=True) val_dataloader, _ = load_data(training=False, cfg=conf, return_dataset=True) train_dataloader = MultiEpochsDataLoader(train_dataloader, batch_size=conf['data']['train_batch_size'], shuffle=True, num_workers=conf['data']['num_workers'], pin_memory=conf['base']['pin_memory'], drop_last=True) val_dataloader = MultiEpochsDataLoader(val_dataloader, batch_size=conf['data']['val_batch_size'], shuffle=False, num_workers=conf['data']['num_workers'], pin_memory=conf['base']['pin_memory'], drop_last=False) # 初始化模型和度量 model, metric = initialize_model_and_metric(conf, class_num) device = conf['base']['device'] model = model.to(device) metric = metric.to(device) # 设置损失函数、优化器和调度器 criterion = setup_loss_function(conf) optimizer, scheduler = setup_optimizer_and_scheduler(conf, model, metric) # 检查点目录 checkpoints = conf['training']['checkpoints'] os.makedirs(checkpoints, exist_ok=True) # GradScaler for mixed precision scaler = torch.cuda.amp.GradScaler() # 更新组件字典 components.update({ 'train_dataloader': train_dataloader, 'val_dataloader': val_dataloader, 'model': model, 'metric': metric, 'criterion': criterion, 'optimizer': optimizer, 'scheduler': scheduler, 'checkpoints': checkpoints, 'scaler': scaler, 'device': device }) return components def run_training_loop(components): """运行完整的训练循环""" # 解包组件 conf = components['conf'] train_dataloader = components['train_dataloader'] val_dataloader = components['val_dataloader'] model = components['model'] metric = components['metric'] criterion = components['criterion'] optimizer = components['optimizer'] scheduler = components['scheduler'] checkpoints = components['checkpoints'] scaler = components['scaler'] device = components['device'] # 训练状态 train_losses = [] val_losses = [] epochs = [] temp_loss = 100 if conf['training']['restore']: print('load pretrain model: {}'.format(conf['training']['restore_model'])) model.load_state_dict(torch.load(conf['training']['restore_model'], map_location=device)) # 训练循环 for e in range(conf['training']['epochs']): train_loss_avg = train_one_epoch(model, metric, criterion, optimizer, train_dataloader, device, scaler, conf) train_losses.append(train_loss_avg) epochs.append(e) val_loss_avg = validate(model, metric, criterion, val_dataloader, device, conf) val_losses.append(val_loss_avg) if val_loss_avg < temp_loss: save_model(model, osp.join(checkpoints, 'best.pth'), isinstance(model, nn.DataParallel)) temp_loss = val_loss_avg scheduler.step() current_lr = optimizer.param_groups[0]['lr'] log_info = ("[{:%Y-%m-%d %H:%M:%S}] Epoch {}/{}, train_loss: {}, val_loss: {} lr:{}" .format(datetime.now(), e, conf['training']['epochs'], train_loss_avg, val_loss_avg, current_lr)) print(log_info) log_training_info(osp.join(conf['logging']['logging_dir']), log_info) print("第%d个epoch的学习率：%f" % (e, current_lr)) # 保存最终模型 save_model(model, osp.join(checkpoints, 'last.pth'), isinstance(model, nn.DataParallel)) # 绘制损失曲线 plt.plot(epochs, train_losses, color='blue', label='Train Loss') plt.plot(epochs, val_losses, color='red', label='Validation Loss') plt.legend() plt.savefig('loss/mobilenetv3Large_2250_0316.png') def main(): """主函数入口""" # 加载配置 conf = load_configuration() # 检查是否启用分布式训练 distributed = conf['base']['distributed'] if distributed: # 分布式训练：使用mp.spawn启动多个进程 local_size = torch.cuda.device_count() world_size = int(conf['distributed']['node_num'])*local_size mp.spawn( run_training, args=(conf['distributed']['node_rank'], local_size, world_size, conf), nprocs=local_size, join=True ) else: # 单机训练：直接运行训练流程 components = initialize_training_components(distributed=False) run_training_loop(components) def run_training(local_rank, node_rank, local_size, world_size, conf): """实际执行训练的函数，供mp.spawn调用""" # 初始化分布式环境 rank = local_rank + node_rank * local_size os.environ['RANK'] = str(rank) os.environ['WORLD_SIZE'] = str(world_size) os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' dist.init_process_group(backend='nccl', rank=rank, world_size=world_size) torch.cuda.set_device(local_rank) device = torch.device('cuda', local_rank) # 获取数据集而不是DataLoader train_dataset, class_num = load_data(training=True, cfg=conf, return_dataset=True) val_dataset, _ = load_data(training=False, cfg=conf, return_dataset=True) # 初始化模型和度量 model, metric = initialize_model_and_metric(conf, class_num) model = model.to(device) metric = metric.to(device) # 包装为DistributedDataParallel模型 model = DDP(model, device_ids=[local_rank], output_device=local_rank) metric = DDP(metric, device_ids=[local_rank], output_device=local_rank) # 设置损失函数、优化器和调度器 criterion = setup_loss_function(conf) optimizer, scheduler = setup_optimizer_and_scheduler(conf, model, metric) # 检查点目录 checkpoints = conf['training']['checkpoints'] os.makedirs(checkpoints, exist_ok=True) # GradScaler for mixed precision scaler = torch.cuda.amp.GradScaler() # 创建分布式采样器 train_sampler = DistributedSampler(train_dataset, shuffle=True) val_sampler = DistributedSampler(val_dataset, shuffle=False) # 使用 MultiEpochsDataLoader 创建分布式数据加载器 train_dataloader = MultiEpochsDataLoader( train_dataset, batch_size=conf['data']['train_batch_size'], sampler=train_sampler, num_workers=conf['data']['num_workers'], pin_memory=conf['base']['pin_memory'], drop_last=True ) val_dataloader = MultiEpochsDataLoader( val_dataset, batch_size=conf['data']['val_batch_size'], sampler=val_sampler, num_workers=conf['data']['num_workers'], pin_memory=conf['base']['pin_memory'], drop_last=False ) # 构建组件字典 components = { 'conf': conf, 'train_dataloader': train_dataloader, 'val_dataloader': val_dataloader, 'model': model, 'metric': metric, 'criterion': criterion, 'optimizer': optimizer, 'scheduler': scheduler, 'checkpoints': checkpoints, 'scaler': scaler, 'device': device, 'distributed': True # 因为是在mp.spawn中运行 } # 运行训练循环 run_training_loop(components) if __name__ == '__main__': main()