import argparse import torch import os from torch import optim from torch.optim import lr_scheduler from util.data_utils import get_data from util.model_utils import get_model from util.train_utils import train_model, validate_model, update_model_weights, v3_update_model_weights def main(args): device = torch.device(args.device) # 数据加载器 loader1, loader2, loader3, subset_len, val_loader = get_data( args.train_path, args.val_path, args.batch_size, args.number_workers ) # 模型 get_model(name='ResNet', number_class=2, device=device, resnet_type='resnet18') model_a = get_model(args.model_name, args.number_class, device, args.deep_backbone).to(device) model_b = get_model(args.model_name, args.number_class, device, args.deep_backbone).to(device) model_c = get_model(args.model_name, args.number_class, device, args.deep_backbone).to(device) # 添加全局模型 global_model = get_model(args.model_name, args.number_class, device, args.deep_backbone).to(device) if args.resume_training: model_a.load_state_dict(torch.load(os.path.join(args.save_dir, 'best_model_a.pth'))) model_b.load_state_dict(torch.load(os.path.join(args.save_dir, 'best_model_b.pth'))) model_c.load_state_dict(torch.load(os.path.join(args.save_dir, 'best_model_c.pth'))) print("已加载之前保存的模型参数继续训练") # 优化器和损失函数 criterion = torch.nn.BCEWithLogitsLoss().to(device) optimizer_a = optim.Adam(model_a.parameters(), lr=args.lr, weight_decay=5e-4) optimizer_b = optim.Adam(model_b.parameters(), lr=args.lr, weight_decay=5e-4) optimizer_c = optim.Adam(model_c.parameters(), lr=args.lr, weight_decay=5e-4) scheduler_a = lr_scheduler.ReduceLROnPlateau(optimizer_a, mode='min', factor=0.5, patience=2, verbose=True) scheduler_b = lr_scheduler.ReduceLROnPlateau(optimizer_b, mode='min', factor=0.5, patience=2, verbose=True) scheduler_c = lr_scheduler.ReduceLROnPlateau(optimizer_c, mode='min', factor=0.5, patience=2, verbose=True) # 初始化最优验证损失和模型路径 best_val_loss_a = float('inf') best_val_loss_b = float('inf') best_val_loss_c = float('inf') save_dir = args.save_dir os.makedirs(save_dir, exist_ok=True) # 训练与验证 for epoch in range(args.epochs): print(f'Epoch {epoch + 1}/{args.epochs}') # 训练模型 loss_a = train_model(device, model_a, loader1, optimizer_a, criterion, epoch, 'model_a') loss_b = train_model(device, model_b, loader2, optimizer_b, criterion, epoch, 'model_b') loss_c = train_model(device, model_c, loader3, optimizer_c, criterion, epoch, 'model_c') # 验证模型 val_loss_a, val_acc_a, val_auc_a = validate_model(device, model_a, val_loader, criterion, epoch, 'model_a') val_loss_b, val_acc_b, val_auc_b = validate_model(device, model_b, val_loader, criterion, epoch, 'model_b') val_loss_c, val_acc_c, val_auc_c = validate_model(device, model_c, val_loader, criterion, epoch, 'model_c') if args.save_model and val_loss_a < best_val_loss_a: best_val_loss_a = val_loss_a torch.save(model_a.state_dict(), os.path.join(save_dir, 'best_model_a.pth')) print(f"Best model_a saved with val_loss: {best_val_loss_a:.4f}") if args.save_model and val_loss_b < best_val_loss_b: best_val_loss_b = val_loss_b torch.save(model_b.state_dict(), os.path.join(save_dir, 'best_model_b.pth')) print(f"Best model_b saved with val_loss: {best_val_loss_b:.4f}") if args.save_model and val_loss_c < best_val_loss_c: best_val_loss_c = val_loss_c torch.save(model_c.state_dict(), os.path.join(save_dir, 'best_model_c.pth')) print(f"Best model_c saved with val_loss: {best_val_loss_c:.4f}") print( f'Model A - Loss: {loss_a:.4f}, Val Loss: {val_loss_a:.4f}, Val Acc: {val_acc_a:.4f}, AUC: {val_auc_a:.4f}') print( f'Model B - Loss: {loss_b:.4f}, Val Loss: {val_loss_b:.4f}, Val Acc: {val_acc_b:.4f}, AUC: {val_auc_b:.4f}') print( f'Model C - Loss: {loss_c:.4f}, Val Loss: {val_loss_c:.4f}, Val Acc: {val_acc_c:.4f}, AUC: {val_auc_c:.4f}') # 更新模型 A 的权重，每 3 轮 1 val_acc_a, val_auc_a, val_acc_a_threshold = v3_update_model_weights( epoch=epoch, model_to_update=model_a, other_models=[model_a, model_b, model_c], global_model=global_model, losses=[loss_a, loss_b, loss_c], val_loader=val_loader, device=device, val_auc_threshold=val_auc_a, validate_model=validate_model, criterion=criterion, update_frequency=1 ) # 更新模型 B 的权重，每 5 轮1 val_acc_b, val_auc_b, val_acc_b_threshold = v3_update_model_weights( epoch=epoch, model_to_update=model_b, other_models=[model_a, model_b, model_c], global_model=global_model, losses=[loss_a, loss_b, loss_c], val_loader=val_loader, device=device, val_auc_threshold=val_auc_b, validate_model=validate_model, criterion=criterion, update_frequency=1 ) # 更新模型 C 的权重，每 2 轮 1 val_acc_c, val_auc_c, val_acc_c_threshold = v3_update_model_weights( epoch=epoch, model_to_update=model_c, other_models=[model_a, model_b, model_c], global_model=global_model, losses=[loss_a, loss_b, loss_c], val_loader=val_loader, device=device, val_auc_threshold=val_auc_c, validate_model=validate_model, criterion=criterion, update_frequency=1 ) print("Training complete! Best models saved.") if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--model_name', type=str, default='resnet18_psa', help='Model name') parser.add_argument('--deep_backbone', type=str, default='*', help='deeplab backbone') parser.add_argument('--train_path', type=str, default='/media/terminator/实验&代码/yhs/FF++/c40/total/train') parser.add_argument('--val_path', type=str, default='/media/terminator/实验&代码/yhs/FF++/c40/total/val') # parser.add_argument('--train_path', type=str, default='/media/terminator/实验&代码/yhs/FF++_mask_sample/c23/df/train') # parser.add_argument('--val_path', type=str, default='/media/terminator/实验&代码/yhs/FF++_mask_sample/c23/df/val') parser.add_argument('--epochs', type=int, default=10) parser.add_argument('--batch_size', type=int, default=16) parser.add_argument('--number_workers', type=int, default=8) parser.add_argument('--number_class', type=int, default=1) parser.add_argument('--device', type=str, default='cuda:0') parser.add_argument('--lr', type=float, default=0.00005) parser.add_argument('--save_dir', type=str, default='/media/terminator/实验&代码/yhs/output/work2/resnet18_psa/c40/total/e10', help='Directory to save best models') parser.add_argument('--save_model', type=bool, default=True, help='是否保存最优模型') parser.add_argument('--resume_training', type=bool, default=False, help='是否从保存的模型参数继续训练') args = parser.parse_args() main(args)