def run(args): with open(args.cfg_path) as f: cfg = edict(json.load(f)) if args.verbose is True: print(json.dumps(cfg, indent=4)) if not os.path.exists(args.save_path): os.mkdir(args.save_path) if args.logtofile is True: logging.basicConfig(filename=args.save_path + '/log.txt', filemode="w", level=logging.INFO) else: logging.basicConfig(level=logging.INFO) if not args.resume: with open(os.path.join(args.save_path, 'cfg.json'), 'w') as f: json.dump(cfg, f, indent=1) device_ids = list(map(int, args.device_ids.split(','))) num_devices = torch.cuda.device_count() if num_devices < len(device_ids): raise Exception('#available gpu : {} < --device_ids : {}'.format( num_devices, len(device_ids))) device = torch.device('cuda:{}'.format(device_ids[0])) model = Classifier(cfg) if args.verbose is True: from torchsummary import summary if cfg.fix_ratio: h, w = cfg.long_side, cfg.long_side else: h, w = cfg.height, cfg.width summary(model.to(device), (3, h, w)) model = DataParallel(model, device_ids=device_ids).to(device).train() if args.pre_train is not None: if os.path.exists(args.pre_train): ckpt = torch.load(args.pre_train, map_location=device) model.module.load_state_dict(ckpt) optimizer = get_optimizer(model.parameters(), cfg) src_folder = os.path.dirname(os.path.abspath(__file__)) + '/../' dst_folder = os.path.join(args.save_path, 'classification') # rc, size = subprocess.getstatusoutput('dir --max-depth=0 %s | cut -f1' # % src_folder) # if rc != 0: # print(size) # raise Exception('Copy folder error : {}'.format(rc)) # rc, err_msg = subprocess.getstatusoutput('cp -R %s %s' % (src_folder, # dst_folder)) # if rc != 0: # raise Exception('copy folder error : {}'.format(err_msg)) copyfile(cfg.train_csv, os.path.join(args.save_path, 'train.csv')) copyfile(cfg.dev_csv, os.path.join(args.save_path, 'valid.csv')) dataloader_train = DataLoader(ImageDataset(cfg.train_csv, cfg, mode='train'), batch_size=cfg.train_batch_size, num_workers=args.num_workers, drop_last=True, shuffle=True) dataloader_dev = DataLoader(ImageDataset(cfg.dev_csv, cfg, mode='dev'), batch_size=cfg.dev_batch_size, num_workers=args.num_workers, drop_last=False, shuffle=False) dev_header = dataloader_dev.dataset._label_header summary_train = {'epoch': 0, 'step': 0} summary_dev = {'loss': float('inf'), 'acc': 0.0} summary_writer = SummaryWriter(args.save_path) epoch_start = 0 best_dict = { "acc_dev_best": 0.0, "auc_dev_best": 0.0, "loss_dev_best": float('inf'), "fused_dev_best": 0.0, "best_idx": 1 } if args.resume: ckpt_path = os.path.join(args.save_path, 'train.ckpt') ckpt = torch.load(ckpt_path, map_location=device) model.module.load_state_dict(ckpt['state_dict']) summary_train = {'epoch': ckpt['epoch'], 'step': ckpt['step']} best_dict['acc_dev_best'] = ckpt['acc_dev_best'] best_dict['loss_dev_best'] = ckpt['loss_dev_best'] best_dict['auc_dev_best'] = ckpt['auc_dev_best'] epoch_start = ckpt['epoch'] for epoch in range(epoch_start, cfg.epoch): lr = lr_schedule(cfg.lr, cfg.lr_factor, summary_train['epoch'], cfg.lr_epochs) for param_group in optimizer.param_groups: param_group['lr'] = lr summary_train, best_dict = train_epoch(summary_train, summary_dev, cfg, args, model, dataloader_train, dataloader_dev, optimizer, summary_writer, best_dict, dev_header) time_now = time.time() summary_dev, predlist, true_list = test_epoch(summary_dev, cfg, args, model, dataloader_dev) time_spent = time.time() - time_now auclist = [] for i in range(len(cfg.num_classes)): y_pred = predlist[i] y_true = true_list[i] fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred, pos_label=1) auc = metrics.auc(fpr, tpr) auclist.append(auc) summary_dev['auc'] = np.array(auclist) loss_dev_str = ' '.join( map(lambda x: '{:.5f}'.format(x), summary_dev['loss'])) acc_dev_str = ' '.join( map(lambda x: '{:.3f}'.format(x), summary_dev['acc'])) auc_dev_str = ' '.join( map(lambda x: '{:.3f}'.format(x), summary_dev['auc'])) logging.info('{}, Dev, Step : {}, Loss : {}, Acc : {}, Auc : {},' 'Mean auc: {:.3f} ' 'Run Time : {:.2f} sec'.format( time.strftime("%Y-%m-%d %H:%M:%S"), summary_train['step'], loss_dev_str, acc_dev_str, auc_dev_str, summary_dev['auc'].mean(), time_spent)) for t in range(len(cfg.num_classes)): summary_writer.add_scalar('dev/loss_{}'.format(dev_header[t]), summary_dev['loss'][t], summary_train['step']) summary_writer.add_scalar('dev/acc_{}'.format(dev_header[t]), summary_dev['acc'][t], summary_train['step']) summary_writer.add_scalar('dev/auc_{}'.format(dev_header[t]), summary_dev['auc'][t], summary_train['step']) save_best = False mean_acc = summary_dev['acc'][cfg.save_index].mean() if mean_acc >= best_dict['acc_dev_best']: best_dict['acc_dev_best'] = mean_acc if cfg.best_target == 'acc': save_best = True mean_auc = summary_dev['auc'][cfg.save_index].mean() if mean_auc >= best_dict['auc_dev_best']: best_dict['auc_dev_best'] = mean_auc if cfg.best_target == 'auc': save_best = True mean_loss = summary_dev['loss'][cfg.save_index].mean() if mean_loss <= best_dict['loss_dev_best']: best_dict['loss_dev_best'] = mean_loss if cfg.best_target == 'loss': save_best = True if save_best: torch.save( { 'epoch': summary_train['epoch'], 'step': summary_train['step'], 'acc_dev_best': best_dict['acc_dev_best'], 'auc_dev_best': best_dict['auc_dev_best'], 'loss_dev_best': best_dict['loss_dev_best'], 'state_dict': model.module.state_dict() }, os.path.join(args.save_path, 'best{}.ckpt'.format(best_dict['best_idx']))) best_dict['best_idx'] += 1 if best_dict['best_idx'] > cfg.save_top_k: best_dict['best_idx'] = 1 logging.info('{}, Best, Step : {}, Loss : {}, Acc : {},' 'Auc :{},Best Auc : {:.3f}'.format( time.strftime("%Y-%m-%d %H:%M:%S"), summary_train['step'], loss_dev_str, acc_dev_str, auc_dev_str, best_dict['auc_dev_best'])) torch.save( { 'epoch': summary_train['epoch'], 'step': summary_train['step'], 'acc_dev_best': best_dict['acc_dev_best'], 'auc_dev_best': best_dict['auc_dev_best'], 'loss_dev_best': best_dict['loss_dev_best'], 'state_dict': model.module.state_dict() }, os.path.join(args.save_path, 'train.ckpt')) summary_writer.close()
def run(args, val_h5_file): with open(args.cfg_path) as f: cfg = edict(json.load(f)) if args.verbose is True: print(json.dumps(cfg, indent=4)) if not os.path.exists(args.save_path): os.mkdir(args.save_path) if args.logtofile is True: logging.basicConfig(filename=args.save_path + '/log.txt', filemode="w", level=logging.INFO) else: logging.basicConfig(level=logging.INFO) if not args.resume: with open(os.path.join(args.save_path, 'cfg.json'), 'w') as f: json.dump(cfg, f, indent=1) device_ids = list(map(int, args.device_ids.split(','))) num_devices = torch.cuda.device_count() if num_devices < len(device_ids): raise Exception('#available gpu : {} < --device_ids : {}'.format( num_devices, len(device_ids))) device = torch.device('cuda:{}'.format(device_ids[0])) model = Classifier(cfg) if args.verbose is True: from torchsummary import summary if cfg.fix_ratio: h, w = cfg.long_side, cfg.long_side else: h, w = cfg.height, cfg.width summary(model.to(device), (3, h, w)) model = DataParallel(model, device_ids=device_ids).to(device).train() if args.pre_train is not None: if os.path.exists(args.pre_train): ckpt = torch.load(args.pre_train, map_location=device) model.module.load_state_dict(ckpt) optimizer = get_optimizer(model.parameters(), cfg) #src_folder = os.path.dirname(os.path.abspath(__file__)) + '/../' #dst_folder = os.path.join(args.save_path, 'classification') #rc, size = subprocess.getstatusoutput('du --max-depth=0 %s | cut -f1' % src_folder) #if rc != 0: raise Exception('Copy folder error : {}'.format(rc)) #rc, err_msg = subprocess.getstatusoutput('cp -R %s %s' % (src_folder, dst_folder)) #if rc != 0: raise Exception('copy folder error : {}'.format(err_msg)) #copyfile(cfg.train_csv, os.path.join(args.save_path, 'train.csv')) #copyfile(cfg.dev_csv, os.path.join(args.save_path, 'dev.csv')) # np_train_h5_file = np.array(train_h5_file['train'][:10000], dtype=np.uint8) # np_t_u_ones = np.array(train_h5_file['train_u_ones'][:10000], dtype=np.int8) # np_t_u_zeros = np.array(train_h5_file['train_u_zeros'][:10000], dtype=np.int8) # np_t_u_random = np.array(train_h5_file['train_u_random'][:10000], dtype=np.int8) np_val_h5_file = np.array(val_h5_file['val'], dtype=np.uint8) np_v_u_ones = np.array(val_h5_file['val_u_ones'], dtype=np.int8) np_v_u_zeros = np.array(val_h5_file['val_u_zeros'], dtype=np.int8) np_v_u_random = np.array(val_h5_file['val_u_random'], dtype=np.int8) train_labels = {} with h5py.File(f'{args.train_chunks}/train_labels.h5', 'r') as fp: train_labels['train_u_ones'] = np.array(fp['train_u_ones'], dtype=np.int8) train_labels['train_u_zeros'] = np.array(fp['train_u_zeros'], dtype=np.int8) train_labels['train_u_random'] = np.array(fp['train_u_random'], dtype=np.int8) np_train_samples = None for i in range(args.chunk_count): with open(f'{args.train_chunks}/chexpert_dset_chunk_{i+1}.npy', 'rb') as f: if np_train_samples is None: np_train_samples = np.load(f) else: np_train_samples = np.concatenate( (np_train_samples, np.load(f))) dataloader_train = DataLoader(ImageDataset( [np_train_samples, train_labels], cfg, mode='train'), batch_size=cfg.train_batch_size, num_workers=args.num_workers, drop_last=True, shuffle=True) dataloader_dev = DataLoader(ImageDataset( [np_val_h5_file, np_v_u_zeros, np_v_u_ones, np_v_u_random], cfg, mode='val'), batch_size=cfg.dev_batch_size, num_workers=args.num_workers, drop_last=False, shuffle=False) #dev_header = dataloader_dev.dataset._label_header dev_header = [ 'No_Finding', 'Enlarged_Cardiomediastinum', 'Cardiomegaly', 'Lung_Opacity', 'Lung_Lesion', 'Edema', 'Consolidation', 'Pneumonia', 'Atelectasis', 'Pneumothorax', 'Pleural_Effusion', 'Pleural_Other', 'Fracture', 'Support_Devices' ] print(f'dataloaders are set. train count: {np_train_samples.shape[0]}') logging.info("[LOGGING TEST]: dataloaders are set...") summary_train = {'epoch': 0, 'step': 0} summary_dev = {'loss': float('inf'), 'acc': 0.0} summary_writer = SummaryWriter(args.save_path) epoch_start = 0 best_dict = { "acc_dev_best": 0.0, "auc_dev_best": 0.0, "loss_dev_best": float('inf'), "fused_dev_best": 0.0, "best_idx": 1 } if args.resume: ckpt_path = os.path.join(args.save_path, 'train.ckpt') ckpt = torch.load(ckpt_path, map_location=device) model.module.load_state_dict(ckpt['state_dict']) summary_train = {'epoch': ckpt['epoch'], 'step': ckpt['step']} best_dict['acc_dev_best'] = ckpt['acc_dev_best'] best_dict['loss_dev_best'] = ckpt['loss_dev_best'] best_dict['auc_dev_best'] = ckpt['auc_dev_best'] epoch_start = ckpt['epoch'] q_list = [] k_list = [] for i in range(len(cfg.num_classes)): q_list.append(args.q) k_list.append(args.k) k_list = torch.FloatTensor(k_list) q_list = torch.FloatTensor(q_list) loss_sq_hinge = MultiClassSquaredHingeLoss() print('Everything is set starting to train...') before = datetime.datetime.now() for epoch in range(epoch_start, cfg.epoch): lr = lr_schedule(cfg.lr, cfg.lr_factor, summary_train['epoch'], cfg.lr_epochs) for param_group in optimizer.param_groups: param_group['lr'] = lr summary_train, best_dict = train_epoch(summary_train, summary_dev, cfg, args, model, dataloader_train, dataloader_dev, optimizer, summary_writer, best_dict, dev_header, q_list, k_list, loss_sq_hinge) time_now = time.time() summary_dev, predlist, true_list = test_epoch(summary_dev, cfg, args, model, dataloader_dev, q_list, k_list, loss_sq_hinge) time_spent = time.time() - time_now auclist = [] for i in range(len(cfg.num_classes)): y_pred = predlist[i] y_true = true_list[i] fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred, pos_label=1) auc = metrics.auc(fpr, tpr) auclist.append(auc) summary_dev['auc'] = np.array(auclist) loss_dev_str = ' '.join( map(lambda x: '{:.5f}'.format(x), summary_dev['loss'])) acc_dev_str = ' '.join( map(lambda x: '{:.3f}'.format(x), summary_dev['acc'])) auc_dev_str = ' '.join( map(lambda x: '{:.3f}'.format(x), summary_dev['auc'])) logging.info('{}, Dev, Step : {}, Loss : {}, Acc : {}, Auc : {},' 'Mean auc: {:.3f} ' 'Run Time : {:.2f} sec'.format( time.strftime("%Y-%m-%d %H:%M:%S"), summary_train['step'], loss_dev_str, acc_dev_str, auc_dev_str, summary_dev['auc'].mean(), time_spent)) for t in range(len(cfg.num_classes)): summary_writer.add_scalar('dev/loss_{}'.format(dev_header[t]), summary_dev['loss'][t], summary_train['step']) summary_writer.add_scalar('dev/acc_{}'.format(dev_header[t]), summary_dev['acc'][t], summary_train['step']) summary_writer.add_scalar('dev/auc_{}'.format(dev_header[t]), summary_dev['auc'][t], summary_train['step']) save_best = False mean_acc = summary_dev['acc'][cfg.save_index].mean() if mean_acc >= best_dict['acc_dev_best']: best_dict['acc_dev_best'] = mean_acc if cfg.best_target == 'acc': save_best = True mean_auc = summary_dev['auc'][cfg.save_index].mean() if mean_auc >= best_dict['auc_dev_best']: best_dict['auc_dev_best'] = mean_auc if cfg.best_target == 'auc': save_best = True mean_loss = summary_dev['loss'][cfg.save_index].mean() if mean_loss <= best_dict['loss_dev_best']: best_dict['loss_dev_best'] = mean_loss if cfg.best_target == 'loss': save_best = True if save_best: torch.save( { 'epoch': summary_train['epoch'], 'step': summary_train['step'], 'acc_dev_best': best_dict['acc_dev_best'], 'auc_dev_best': best_dict['auc_dev_best'], 'loss_dev_best': best_dict['loss_dev_best'], 'state_dict': model.module.state_dict() }, os.path.join(args.save_path, 'best{}.ckpt'.format(best_dict['best_idx']))) best_dict['best_idx'] += 1 if best_dict['best_idx'] > cfg.save_top_k: best_dict['best_idx'] = 1 logging.info('{}, Best, Step : {}, Loss : {}, Acc : {},' 'Auc :{},Best Auc : {:.3f}'.format( time.strftime("%Y-%m-%d %H:%M:%S"), summary_train['step'], loss_dev_str, acc_dev_str, auc_dev_str, best_dict['auc_dev_best'])) torch.save( { 'epoch': summary_train['epoch'], 'step': summary_train['step'], 'acc_dev_best': best_dict['acc_dev_best'], 'auc_dev_best': best_dict['auc_dev_best'], 'loss_dev_best': best_dict['loss_dev_best'], 'state_dict': model.module.state_dict() }, os.path.join(args.save_path, 'train.ckpt')) print_remaining_time(before, epoch + 1, cfg.epoch, additional='[training]') summary_writer.close()
def main(): if args.dataset == 'ChestXray-NIHCC': if args.no_fiding: classes = [ 'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening', 'Hernia', 'No Fiding' ] else: classes = [ 'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening', 'Hernia' ] elif args.dataset == 'CheXpert-v1.0-small': classes = [ 'No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung Opacity', 'Lung Lesion', 'Edema', 'Consolidation', 'Pneumonia', 'Atelectasis', 'Pneumothorax', 'Pleural Effusion', 'Pleural Other', 'Fracture', 'Support Devices' ] else: print('--dataset incorrect') return torch.manual_seed(args.seed) use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") pin_memory = True if use_gpu else False print("Initializing dataset: {}".format(args.dataset)) data_transforms = { 'train': transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.Resize(556), transforms.CenterCrop(512), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]), 'valid': transforms.Compose([ transforms.Resize(556), transforms.CenterCrop(512), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]), } datasetTrain = DatasetGenerator(path_base=args.base_dir, dataset_file='train', transform=data_transforms['train'], dataset_=args.dataset, no_fiding=args.no_fiding) datasetVal = DatasetGenerator(path_base=args.base_dir, dataset_file='valid', transform=data_transforms['valid'], dataset_=args.dataset, no_fiding=args.no_fiding) train_loader = DataLoader(dataset=datasetTrain, batch_size=args.train_batch, shuffle=args.train_shuffle, num_workers=args.workers, pin_memory=pin_memory) valid_loader = DataLoader(dataset=datasetVal, batch_size=args.valid_batch, shuffle=args.valid_shuffle, num_workers=args.workers, pin_memory=pin_memory) with open(args.infos_densenet) as f: cfg = edict(json.load(f)) print('Initializing densenet branch') model_dense = Classifier(cfg) print("Model size: {:.5f}M".format( sum(p.numel() for p in model_dense.parameters()) / 1000000.0)) with open(args.infos_resnet) as f: cfg = edict(json.load(f)) print('Initializing resnet branch') model_res = Classifier(cfg) print("Model size: {:.5f}M".format( sum(p.numel() for p in model_res.parameters()) / 1000000.0)) print('Initializing fusion branch') model_fusion = Fusion(input_size=7424, output_size=len(classes)) print("Model size: {:.5f}M".format( sum(p.numel() for p in model_fusion.parameters()) / 1000000.0)) print("Initializing optimizers") optimizer_dense = init_optim(args.optim, model_dense.parameters(), args.learning_rate, args.weight_decay, args.momentum) optimizer_res = init_optim(args.optim, model_res.parameters(), args.learning_rate, args.weight_decay, args.momentum) optimizer_fusion = init_optim(args.optim, model_fusion.parameters(), args.learning_rate, args.weight_decay, args.momentum) criterion = nn.BCELoss() print("Initializing scheduler: {}".format(args.scheduler)) if args.stepsize > 0: scheduler_dense = init_scheduler(args.scheduler, optimizer_dense, args.stepsize, args.gamma) scheduler_res = init_scheduler(args.scheduler, optimizer_res, args.stepsize, args.gamma) scheduler_fusion = init_scheduler(args.scheduler, optimizer_fusion, args.stepsize, args.gamma) start_epoch = args.start_epoch best_loss = np.inf if args.resume_densenet: checkpoint_dense = torch.load(args.resume_densenet) model_dense.load_state_dict(checkpoint_dense['state_dict']) epoch_dense = checkpoint_dense['epoch'] print("Resuming densenet from epoch {}".format(epoch_dense + 1)) if args.resume_resnet: checkpoint_res = torch.load(args.resume_resnet) model_res.load_state_dict(checkpoint_res['state_dict']) epoch_res = checkpoint_res['epoch'] print("Resuming resnet from epoch {}".format(epoch_res + 1)) if args.resume_fusion: checkpoint_fusion = torch.load(args.resume_fusion) model_fusion.load_state_dict(checkpoint_fusion['state_dict']) epoch_fusion = checkpoint_fusion['epoch'] print("Resuming fusion from epoch {}".format(epoch_fusion + 1)) if use_gpu: model_dense = nn.DataParallel(model_dense).cuda() model_res = nn.DataParallel(model_res).cuda() model_fusion = nn.DataParallel(model_fusion).cuda() if args.evaluate: print("Evaluate only") if args.step == 1: valid('step1', model_dense, model_res, model_fusion, valid_loader, criterion, args.print_freq, classes, cfg, data_transforms['valid']) elif args.step == 2: valid('step2', model_dense, model_res, model_fusion, valid_loader, criterion, args.print_freq, classes, cfg, data_transforms['valid']) elif args.step == 3: valid('step3', model_dense, model_res, model_fusion, valid_loader, criterion, args.print_freq, classes, cfg, data_transforms['valid']) else: print('args.step not found') return if args.step == 1: #################################### DENSENET BRANCH INIT ########################################## start_time = time.time() train_time = 0 best_epoch = 0 print("==> Start training of densenet branch") for p in model_dense.parameters(): p.requires_grad = True for p in model_res.parameters(): p.requires_grad = False for p in model_fusion.parameters(): p.requires_grad = True for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train('step1', model_dense, model_res, model_fusion, train_loader, optimizer_dense, optimizer_res, optimizer_fusion, criterion, args.print_freq, epoch, args.max_epoch, cfg, data_transforms['train']) train_time += round(time.time() - start_train_time) if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: print("==> Validation") loss_val = valid('step1', model_dense, model_res, model_fusion, valid_loader, criterion, args.print_freq, classes, cfg, data_transforms['valid']) if args.stepsize > 0: if args.scheduler == 'ReduceLROnPlateau': scheduler_dense.step(loss_val) scheduler_fusion.step(loss_val) else: scheduler_dense.step() scheduler_fusion.step() is_best = loss_val < best_loss if is_best: best_loss = loss_val best_epoch = epoch + 1 if use_gpu: state_dict_dense = model_dense.module.state_dict() state_dict_fusion = model_fusion.module.state_dict() else: state_dict_dense = model_dense.state_dict() state_dict_fusion = model_fusion.state_dict() save_checkpoint( { 'state_dict': state_dict_dense, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'dense') save_checkpoint( { 'state_dict': state_dict_fusion, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'fusion') print("==> Best Validation Loss {:.4%}, achieved at epoch {}".format( best_loss, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Dense branch finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}." .format(elapsed, train_time)) #################################### DENSENET BRANCH END ########################################## elif args.step == 2: #################################### RESNET BRANCH INIT ########################################## start_time = time.time() train_time = 0 best_epoch = 0 print("==> Start training of local branch") for p in model_dense.parameters(): p.requires_grad = False for p in model_res.parameters(): p.requires_grad = True for p in model_fusion.parameters(): p.requires_grad = True for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train('step2', model_dense, model_res, model_fusion, train_loader, optimizer_dense, optimizer_res, optimizer_fusion, criterion, args.print_freq, epoch, args.max_epoch, cfg, data_transforms['train']) train_time += round(time.time() - start_train_time) if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: print("==> Validation") loss_val = valid('step2', model_dense, model_res, model_fusion, valid_loader, criterion, args.print_freq, classes, cfg, data_transforms['valid']) if args.stepsize > 0: if args.scheduler == 'ReduceLROnPlateau': scheduler_res.step(loss_val) scheduler_fusion.step(loss_val) else: scheduler_res.step() scheduler_fusion.step() is_best = loss_val < best_loss if is_best: best_loss = loss_val best_epoch = epoch + 1 if use_gpu: state_dict_res = model_res.module.state_dict() state_dict_fusion = model_fusion.module.state_dict() else: state_dict_res = model_res.state_dict() state_dict_fusion = model_fusion.state_dict() save_checkpoint( { 'state_dict': state_dict_res, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'res') save_checkpoint( { 'state_dict': state_dict_fusion, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'fusion') print("==> Best Validation Loss {:.4%}, achieved at epoch {}".format( best_loss, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Resnet branch finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}." .format(elapsed, train_time)) #################################### RESNET BRANCH END ########################################## elif args.step == 3: #################################### FUSION BRANCH INIT ########################################## start_time = time.time() train_time = 0 best_epoch = 0 print("==> Start training of fusion branch") for p in model_dense.parameters(): p.requires_grad = True for p in model_res.parameters(): p.requires_grad = True for p in model_fusion.parameters(): p.requires_grad = True for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train('step3', model_dense, model_res, model_fusion, train_loader, optimizer_dense, optimizer_res, optimizer_fusion, criterion, args.print_freq, epoch, args.max_epoch, cfg, data_transforms['train']) train_time += round(time.time() - start_train_time) if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: print("==> Validation") loss_val = valid('step3', model_dense, model_res, model_fusion, valid_loader, criterion, args.print_freq, classes, cfg, data_transforms['valid']) if args.stepsize > 0: if args.scheduler == 'ReduceLROnPlateau': scheduler_dense.step(loss_val) scheduler_res.step(loss_val) scheduler_fusion.step(loss_val) else: scheduler_dense.step() scheduler_res.step() scheduler_fusion.step() is_best = loss_val < best_loss if is_best: best_loss = loss_val best_epoch = epoch + 1 if use_gpu: state_dict_dense = model_dense.module.state_dict() state_dict_res = model_res.module.state_dict() state_dict_fusion = model_fusion.module.state_dict() else: state_dict_dense = model_dense.state_dict() state_dict_res = model_res.state_dict() state_dict_fusion = model_fusion.state_dict() save_checkpoint( { 'state_dict': state_dict_dense, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'dense') save_checkpoint( { 'state_dict': state_dict_res, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'res') save_checkpoint( { 'state_dict': state_dict_fusion, 'loss': best_loss, 'epoch': epoch, }, is_best, args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar', 'fusion') print("==> Best Validation Loss {:.4%}, achieved at epoch {}".format( best_loss, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Fusion branch finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}." .format(elapsed, train_time)) #################################### FUSION BRANCH END ########################################## else: print('args.step not found')
classifier = Classifier(opts.latent_size).to(device) classer = CLASSIFIERS().to(device) print(cvae) print(dis) print(classifier) optimizer_cvae = torch.optim.Adam(cvae.parameters(), lr=opts.lr, betas=(opts.b1, opts.b2), weight_decay=opts.weight_decay) optimizer_dis = torch.optim.Adam(dis.parameters(), lr=opts.lr, betas=(opts.b1, opts.b2), weight_decay=opts.weight_decay) optimizer_classifier = torch.optim.Adam(classifier.parameters(), lr=opts.lr, betas=(opts.b1, opts.b2), weight_decay=opts.weight_decay) i = 1 while os.path.isdir('./ex/' + str(i)): i += 1 os.mkdir('./ex/' + str(i)) output_path = './ex/' + str(i) losses = { 'total': [], 'kl': [], 'bce': [], 'dis': [],
def run_fl(args): with open(args.cfg_path) as f: cfg = edict(json.load(f)) if args.verbose is True: print(json.dumps(cfg, indent=4)) if not os.path.exists(args.save_path): os.mkdir(args.save_path) if args.logtofile is True: logging.basicConfig(filename=args.save_path + '/log.txt', filemode="w", level=logging.INFO) else: logging.basicConfig(level=logging.INFO) if not args.resume: with open(os.path.join(args.save_path, 'cfg.json'), 'w') as f: json.dump(cfg, f, indent=1) device_ids = list(map(int, args.device_ids.split(','))) num_devices = torch.cuda.device_count() if num_devices < len(device_ids): raise Exception('#available gpu : {} < --device_ids : {}'.format( num_devices, len(device_ids))) device = torch.device('cuda:{}'.format(device_ids[0])) # initialise global model model = Classifier(cfg).to(device).train() if args.verbose is True: from torchsummary import summary if cfg.fix_ratio: h, w = cfg.long_side, cfg.long_side else: h, w = cfg.height, cfg.width summary(model.to(device), (3, h, w)) if args.pre_train is not None: if os.path.exists(args.pre_train): ckpt = torch.load(args.pre_train, map_location=device) model.load_state_dict(ckpt) src_folder = os.path.dirname(os.path.abspath(__file__)) + '/../' dst_folder = os.path.join(args.save_path, 'classification') rc, size = subprocess.getstatusoutput('du --max-depth=0 %s | cut -f1' % src_folder) if rc != 0: raise Exception('Copy folder error : {}'.format(rc)) else: print('Successfully determined size of directory') rc, err_msg = subprocess.getstatusoutput('cp -R %s %s' % (src_folder, dst_folder)) if rc != 0: raise Exception('copy folder error : {}'.format(err_msg)) else: print('Successfully copied folder') # copy train files train_files = cfg.train_csv clients = {} for i, c in enumerate(string.ascii_uppercase): if i < len(train_files): clients[c] = {} else: break # initialise clients for i, client in enumerate(clients): copyfile(train_files[i], os.path.join(args.save_path, f'train_{client}.csv')) clients[client]['dataloader_train'] =\ DataLoader( ImageDataset(train_files[i], cfg, mode='train'), batch_size=cfg.train_batch_size, num_workers=args.num_workers,drop_last=True, shuffle=True ) clients[client]['bytes_uploaded'] = 0.0 clients[client]['epoch'] = 0 copyfile(cfg.dev_csv, os.path.join(args.save_path, 'dev.csv')) dataloader_dev = DataLoader(ImageDataset(cfg.dev_csv, cfg, mode='dev'), batch_size=cfg.dev_batch_size, num_workers=args.num_workers, drop_last=False, shuffle=False) dev_header = dataloader_dev.dataset._label_header w_global = model.state_dict() summary_train = {'epoch': 0, 'step': 0} summary_dev = {'loss': float('inf'), 'acc': 0.0} summary_writer = SummaryWriter(args.save_path) comm_rounds = cfg.epoch best_dict = { "acc_dev_best": 0.0, "auc_dev_best": 0.0, "loss_dev_best": float('inf'), "fused_dev_best": 0.0, "best_idx": 1 } # Communication rounds loop for cr in range(comm_rounds): logging.info('{}, Start communication round {} of FL - {} ...'.format( time.strftime("%Y-%m-%d %H:%M:%S"), cr + 1, cfg.fl_technique)) w_locals = [] for client in clients: logging.info( '{}, Start local training process for client {}, communication round: {} ...' .format(time.strftime("%Y-%m-%d %H:%M:%S"), client, cr + 1)) # Load previous current global model as start point model = Classifier(cfg).to(device).train() model.load_state_dict(w_global) if cfg.fl_technique == "FedProx": global_weight_collector = get_global_weights(model, device) else: global_weight_collector = None optimizer = get_optimizer(model.parameters(), cfg) # local training loops for epoch in range(cfg.local_epoch): lr = lr_schedule(cfg.lr, cfg.lr_factor, epoch, cfg.lr_epochs) for param_group in optimizer.param_groups: param_group['lr'] = lr summary_train, best_dict = train_epoch_fl( summary_train, summary_dev, cfg, args, model, clients[client]['dataloader_train'], dataloader_dev, optimizer, summary_writer, best_dict, dev_header, epoch, global_weight_collector) summary_train['step'] += 1 bytes_to_upload = sys.getsizeof(model.state_dict()) clients[client]['bytes_uploaded'] += bytes_to_upload logging.info( '{}, Completed local rounds for client {} in communication round {}. ' 'Uploading {} bytes to server, {} bytes in total sent from client' .format(time.strftime("%Y-%m-%d %H:%M:%S"), client, cr + 1, bytes_to_upload, clients[client]['bytes_uploaded'])) w_locals.append(model.state_dict()) if cfg.fl_technique == "FedAvg": w_global = fed_avg(w_locals) elif cfg.fl_technique == 'WFedAvg': w_global = weighted_fed_avg(w_locals, cfg.train_proportions) elif cfg.fl_technique == 'FedProx': # Use weighted FedAvg when using FedProx w_global = weighted_fed_avg(w_locals, cfg.train_proportions) # Test the performance of the averaged model avged_model = Classifier(cfg).to(device) avged_model.load_state_dict(w_global) time_now = time.time() summary_dev, predlist, true_list = test_epoch(summary_dev, cfg, args, avged_model, dataloader_dev) time_spent = time.time() - time_now auclist = [] for i in range(len(cfg.num_classes)): y_pred = predlist[i] y_true = true_list[i] fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred, pos_label=1) auc = metrics.auc(fpr, tpr) auclist.append(auc) auc_summary = np.array(auclist) loss_dev_str = ' '.join( map(lambda x: '{:.5f}'.format(x), summary_dev['loss'])) acc_dev_str = ' '.join( map(lambda x: '{:.3f}'.format(x), summary_dev['acc'])) auc_dev_str = ' '.join(map(lambda x: '{:.3f}'.format(x), auc_summary)) logging.info( '{}, Averaged Model -> Dev, Step : {}, Loss : {}, Acc : {}, Auc : {},' 'Mean auc: {:.3f} ' 'Run Time : {:.2f} sec'.format(time.strftime("%Y-%m-%d %H:%M:%S"), summary_train['step'], loss_dev_str, acc_dev_str, auc_dev_str, auc_summary.mean(), time_spent)) for t in range(len(cfg.num_classes)): summary_writer.add_scalar('dev/loss_{}'.format(dev_header[t]), summary_dev['loss'][t], summary_train['step']) summary_writer.add_scalar('dev/acc_{}'.format(dev_header[t]), summary_dev['acc'][t], summary_train['step']) summary_writer.add_scalar('dev/auc_{}'.format(dev_header[t]), auc_summary[t], summary_train['step']) save_best = False mean_acc = summary_dev['acc'][cfg.save_index].mean() if mean_acc >= best_dict['acc_dev_best']: best_dict['acc_dev_best'] = mean_acc if cfg.best_target == 'acc': save_best = True mean_auc = auc_summary[cfg.save_index].mean() if mean_auc >= best_dict['auc_dev_best']: best_dict['auc_dev_best'] = mean_auc if cfg.best_target == 'auc': save_best = True mean_loss = summary_dev['loss'][cfg.save_index].mean() if mean_loss <= best_dict['loss_dev_best']: best_dict['loss_dev_best'] = mean_loss if cfg.best_target == 'loss': save_best = True if save_best: torch.save( { 'epoch': summary_train['epoch'], 'step': summary_train['step'], 'acc_dev_best': best_dict['acc_dev_best'], 'auc_dev_best': best_dict['auc_dev_best'], 'loss_dev_best': best_dict['loss_dev_best'], 'state_dict': avged_model.state_dict() }, os.path.join(args.save_path, 'best{}.ckpt'.format(best_dict['best_idx']))) best_dict['best_idx'] += 1 if best_dict['best_idx'] > cfg.save_top_k: best_dict['best_idx'] = 1 logging.info('{}, Best, Step : {}, Loss : {}, Acc : {},' 'Auc :{},Best Auc : {:.3f}'.format( time.strftime("%Y-%m-%d %H:%M:%S"), summary_train['step'], loss_dev_str, acc_dev_str, auc_dev_str, best_dict['auc_dev_best'])) torch.save( { 'epoch': cr, 'step': summary_train['step'], 'acc_dev_best': best_dict['acc_dev_best'], 'auc_dev_best': best_dict['auc_dev_best'], 'loss_dev_best': best_dict['loss_dev_best'], 'state_dict': avged_model.state_dict() }, os.path.join(args.save_path, 'train.ckpt'))
device = torch.device(f"cuda:{device_ids[0]}") model = Classifier(cfg) if args.verbose: from torchsummary import summary h, w = (cfg.long_side, cfg.long_side) if cfg.fix_ratio \ else (cfg.height, cfg.width) summary(model.to(device), (3, h, w)) model = DataParallel(model, device_ids=device_ids).to(device) if args.pre_train is not None: if exists(args.pre_train): ckpt = torch.load(args.pre_train, map_location=device) model.module.load_state_dict(ckpt) optimizer = get_optimizer(model.parameters(), cfg) trainset = ImageDataset(cfg.train_csv, cfg, mode='train') testset = ImageDataset(cfg.dev_csv, cfg, mode='val') trainloader = DataLoader(trainset, batch_size=cfg.train_batch_size, num_workers=args.num_workers, drop_last=True, shuffle=True) testloader = DataLoader(testset, batch_size=cfg.dev_batch_size, num_workers=args.num_workers, drop_last=False, shuffle=False) dev_header = testloader.dataset._label_header # Initialize parameters to log training output summary_train = {'epoch': 0, 'step': 0} summary_dev = {'loss': float('inf'), 'acc': 0.0}