def main(args): ############ init config ################ model_name = args.model assert model_name in models_dict.keys(),"The Usage model is not exist !" print('Usage model :{}'.format(model_name)) #################### init logger ################################### log_dir = './logs/'+ args.model+'_'+args.note + '/{}'.format(time.strftime('%Y%m%d-%H%M%S')) logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Train'.format(args.model)) # setting setting={k: v for k, v in args._get_kwargs()} logger.info(setting) args.save_path = log_dir args.save_tbx_log = args.save_path + '/tbx_log' writer = SummaryWriter(args.save_tbx_log) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda= args.gpus>0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### train_loader=get_dataloder(args,split_flag="train") val_loader=get_dataloder(args,split_flag="valid") ######################## init model ############################################ # model logger.info("Model Dict has keys: \n {}".format(models_dict.keys())) model=get_models(args) if torch.cuda.device_count() > 1 and args.use_cuda: logger.info('use: %d gpus', torch.cuda.device_count()) model = nn.DataParallel(model) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss=='bce': criterion=nn.BCELoss() elif args.loss=='bcelog': criterion=nn.BCEWithLogitsLoss() elif args.loss=="dice": criterion=DiceLoss() elif args.loss=="softdice": criterion=SoftDiceLoss() elif args.loss=='bcedice': criterion=BCEDiceLoss() else: criterion=nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model=model.to(args.device) criterion=criterion.to(args.device) # init optimizer if args.model_optimizer=="sgd": #torch.optim.SGD(parametetrs,lr=args.lr,weight_decay=args.weight_decay,momentum=args.momentum) optimizer=torch.optim.SGD(model.parameters(),lr=args.lr,weight_decay=args.weight_decay,momentum=args.momentum) else: optimizer=torch.optim.Adam(model.parameters(),args.lr,[args.beta1, args.beta2], weight_decay=args.weight_decay) # init schedulers Steplr scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,args.epoch) # scheduler=torch.optim.lr_scheduler.StepLR(optimizer=optimizer,step_size=30,gamma=0.1,last_epoch=-1) ############################### check resume ######################### start_epoch=0 if args.resume is not None: if os.path.isfile(args.resume): logger.info("Loading model and optimizer from checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume, map_location=args.device) start_epoch = checkpoint['epoch'] optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['state_dict']) scheduler.load_state_dict(checkpoint['scheduler']) else: raise FileNotFoundError("No checkpoint found at '{}'".format(args.resume)) #################################### train and val ######################## max_value=0 for epoch in range(start_epoch,args.epoch): # lr=adjust_learning_rate(args,optimizer,epoch) scheduler.step() logger.info('Epoch: %d lr %e', epoch, scheduler.get_lr()[0]) # train mr, ms, mp, mf, mjc, md, macc, mean_loss=train(args, model, criterion, train_loader, optimizer, epoch, logger) # write writer.add_scalar('Train/Loss', mean_loss, epoch) writer.add_scalar('Train/mAcc', macc, epoch) writer.add_scalar('Train/Recall', mr, epoch) writer.add_scalar('Train/Specifi', ms, epoch) writer.add_scalar('Train/Precision', mp, epoch) writer.add_scalar('Train/F1', mf, epoch) writer.add_scalar('Train/Jc', mjc, epoch) writer.add_scalar('Train/Dice', md, epoch) # val vmr, vms, vmp, vmf, vmjc, vmd, vmacc, vmean_loss=val(args, model, criterion, val_loader, epoch, logger) writer.add_scalar('Val/Loss', vmean_loss, epoch) writer.add_scalar('Val/mAcc', vmacc, epoch) writer.add_scalar('Val/Recall', vmr, epoch) writer.add_scalar('Val/Specifi', vms, epoch) writer.add_scalar('Val/Precision', vmp, epoch) writer.add_scalar('Val/F1', vmf, epoch) writer.add_scalar('Val/Jc', vmjc, epoch) writer.add_scalar('Val/Dice', vmd, epoch) is_best=True if (vmjc>=max_value) else False max_value=max(max_value,vmjc) state={ 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'state_dict': model.state_dict(), 'scheduler': model.state_dict(), } logger.info("epoch:{} best:{} max_value:{}".format(epoch,is_best,max_value)) if not is_best: torch.save(state,os.path.join(args.save_path,"checkpoint.pth.tar")) else: torch.save(state,os.path.join(args.save_path,"checkpoint.pth.tar")) torch.save(state,os.path.join(args.save_path,"model_best.pth.tar")) writer.close()
def main(args): #################### init logger ################################### args.model='unet' model_weight_path='../logs/isic2018/unet_ep300/20200402-135108/model_best.pth.tar' model=get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) log_dir = './models/' + args.model+'_prune_'+args.note logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-L1Prune'.format(args.model)) # setting args.save_path = log_dir args.save_tbx_log = args.save_path + '/tbx_log' writer = SummaryWriter(args.save_tbx_log) if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) train_loader=get_dataloder(args,split_flag="train") val_loader=get_dataloder(args,split_flag="valid") # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) logger.info("Original trained model performance test: ") infer(args, model, criterion, val_loader,logger) # Pruning # Pruning Configuration, in paper 'PRUNING FILTERS FOR EFFICIENT CONVNETS', configure_list = [{ 'sparsity': 0.5, 'op_types': ['Conv2d'], 'op_names': ['Conv1.conv.0','Conv1.conv.3','Conv2.conv.0','Conv2.conv.3','Conv3.conv.0','Conv3.conv.3', 'Conv4.conv.0','Conv4.conv.3','Conv5.conv.0','Conv5.conv.3', 'Up5.up.1','Up_conv5.conv.0','Up_conv5.conv.3', 'Up4.up.1','Up_conv4.conv.0','Up_conv4.conv.3', 'Up3.up.1','Up_conv3.conv.0','Up_conv3.conv.3', 'Up2.up.1','Up_conv2.conv.0','Up_conv2.conv.3', ]} ] # Prune model and test accuracy without fine tuning. logger.info('=' * 10 + 'Test on the pruned model before fine tune' + '=' * 10) pruner = L1FilterPruner(model, configure_list) # change the forward func (mul pruning mask ) model = pruner.compress() # test performance without finetuning logger.info("Pruning trained model performance test: ") infer(args, model, criterion, val_loader,logger) # Fine tune the pruned model for 40 epochs and test accuracy logger.info('=' * 10 + 'Fine tuning' + '=' * 10) #torch.optim.SGD(parametetrs,lr=args.lr,weight_decay=args.weight_decay,momentum=args.momentum) optimizer=torch.optim.SGD(model.parameters(),lr=args.lr,weight_decay=args.weight_decay,momentum=args.momentum) # init schedulers Steplr scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,args.epoch) max_value = 0 for epoch in range(0, args.epoch): # lr=adjust_learning_rate(args,optimizer,epoch) scheduler.step() logger.info('Epoch: %d lr %e', epoch, scheduler.get_lr()[0]) # update mask pruner.update_epoch(epoch) # train train(args, model, criterion, train_loader,optimizer, epoch, logger) # val vmr, vms, vmp, vmf, vmjc, vmd, vmacc,vloss = infer(args, model, criterion, val_loader,logger) writer.add_scalar('Val/Loss', vloss, epoch) writer.add_scalar('Val/mAcc', vmacc, epoch) writer.add_scalar('Val/Recall', vmr, epoch) writer.add_scalar('Val/Specifi', vms, epoch) writer.add_scalar('Val/Precision', vmp, epoch) writer.add_scalar('Val/F1', vmf, epoch) writer.add_scalar('Val/Jc', vmjc, epoch) writer.add_scalar('Val/Dice', vmd, epoch) is_best = True if (vmjc >= max_value) else False max_value = max(max_value, vmjc) if is_best: pruner.export_model(model_path=os.path.join(args.save_path,"best_prune_unet.pth"), mask_path=os.path.join(args.save_path,'mask_prune_indexs.pth')) state = { 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'state_dict': model.state_dict(), 'scheduler': model.state_dict(), } logger.info("epoch:{} best:{} max_value:{}".format(epoch, is_best, max_value)) torch.save(state, os.path.join(args.save_path, "checkpoint.pth.tar")) writer.close() # test the best_prune_unet.pth args.model='unet' model_weight_path=os.path.join(args.save_path,"best_prune_unet.pth") model=get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')) model = model.to(args.device) logger.info("Final saved pruned model performance test: ") infer(args, model, criterion, val_loader,logger)
def main(args): ############ init config ################ #################### init logger ################################### log_dir = './search_exp/' + '/{}'.format(args.model) + \ '/{}'.format(args.dataset) + '/{}_{}'.format(time.strftime('%Y%m%d-%H%M%S'),args.note) logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Search'.format(args.model)) args.save_path = log_dir args.save_tbx_log = args.save_path + '/tbx_log' writer = SummaryWriter(args.save_tbx_log) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.multi_gpu = args.gpus > 1 and torch.cuda.is_available() args.device = torch.device('cuda:0' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.enabled = True cudnn.benchmark = True setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) ####################### init dataset ########################################### logger.info("Dataset for search is {}".format(args.dataset)) train_dataset = datasets_dict[args.dataset](args, args.dataset_root, split='train') val_dataset = datasets_dict[args.dataset](args, args.dataset_root, split='valid') # train_dataset=datasets_dict[args.dataset](args,split='train') # val_dataset=datasets_dict[args.dataset](args,split='valid') num_train = len(train_dataset) indices = list(range(num_train)) split = int(np.floor(args.train_portion * num_train)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load criterion to gpu !") criterion = criterion.to(args.device) ######################## init model ############################################ switches_normal = [] switches_down = [] switches_up = [] nums_mixop = sum([2 + i for i in range(args.meta_node_num)]) for i in range(nums_mixop): switches_normal.append([True for j in range(len(CellPos))]) for i in range(nums_mixop): switches_down.append([True for j in range(len(CellLinkDownPos))]) for i in range(nums_mixop): switches_up.append([True for j in range(len(CellLinkUpPos))]) # 6-->3-->1 drop_op_down = [2, 3] # 4-->2-->1 drop_op_up = [2, 1] # 7-->4-->1 drop_op_normal = [3, 3] # stage0 pruning stage 1 pruning, stage 2 (training) original_train_batch = args.train_batch original_val_batch = args.val_batch for sp in range(2): # build dataloader # model ,numclass=1,im_ch=3,init_channel=16,intermediate_nodes=4,layers=9 if sp == 0: args.model = "UnetLayer7" args.layers = 7 sp_train_batch = original_train_batch sp_val_batch = original_val_batch sp_epoch = args.epochs sp_lr = args.lr else: #args.model = "UnetLayer9" # 在算力平台上面UnetLayer9就是UnetLayer9_v2 args.model = "UnetLayer9" args.layers = 9 sp_train_batch = original_train_batch sp_val_batch = original_val_batch sp_lr = args.lr sp_epoch = args.epochs train_queue = data.DataLoader( train_dataset, batch_size=sp_train_batch, sampler=torch.utils.data.sampler.SubsetRandomSampler( indices[:split]), pin_memory=True, num_workers=args.num_workers) val_queue = data.DataLoader( train_dataset, batch_size=sp_train_batch, sampler=torch.utils.data.sampler.SubsetRandomSampler( indices[split:num_train]), pin_memory=True, num_workers=args.num_workers) test_dataloader = data.DataLoader(val_dataset, batch_size=sp_val_batch, pin_memory=True, num_workers=args.num_workers) logger.info( "stage:{} model:{} epoch:{} lr:{} train_batch:{} val_batch:{}". format(sp, args.model, sp_epoch, sp_lr, sp_train_batch, sp_val_batch)) model = get_models(args, switches_normal, switches_down, switches_up) save_model_path = os.path.join(args.save_path, "stage_{}_model".format(sp)) if not os.path.exists(save_model_path): os.mkdir(save_model_path) if args.multi_gpu: logger.info('use: %d gpus', args.gpus) model = nn.DataParallel(model) model = model.to(args.device) logger.info('param size = %fMB', calc_parameters_count(model)) # init optimizer for arch parameters and weight parameters # final stage, just train the network parameters optimizer_arch = torch.optim.Adam(model.arch_parameters(), lr=args.arch_lr, betas=(0.5, 0.999), weight_decay=args.arch_weight_decay) optimizer_weight = torch.optim.SGD(model.weight_parameters(), lr=sp_lr, weight_decay=args.weight_decay, momentum=args.momentum) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer_weight, sp_epoch, eta_min=args.lr_min) #################################### train and val ######################## max_value = 0 for epoch in range(0, sp_epoch): # lr=adjust_learning_rate(args,optimizer,epoch) scheduler.step() logger.info('Epoch: %d lr %e', epoch, scheduler.get_lr()[0]) # train if epoch < args.arch_after: weight_loss_avg, arch_loss_avg, mr, ms, mp, mf, mjc, md, macc = train( args, train_queue, val_queue, model, criterion, optimizer_weight, optimizer_arch, train_arch=False) else: weight_loss_avg, arch_loss_avg, mr, ms, mp, mf, mjc, md, macc = train( args, train_queue, val_queue, model, criterion, optimizer_weight, optimizer_arch, train_arch=True) logger.info("Epoch:{} WeightLoss:{:.3f} ArchLoss:{:.3f}".format( epoch, weight_loss_avg, arch_loss_avg)) logger.info(" Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}".format( macc, md, mjc)) # write writer.add_scalar('Train/W_loss', weight_loss_avg, epoch) writer.add_scalar('Train/A_loss', arch_loss_avg, epoch) writer.add_scalar('Train/Dice', md, epoch) # infer if (epoch + 1) % args.infer_epoch == 0: genotype = model.genotype() logger.info('genotype = %s', genotype) val_loss, (vmr, vms, vmp, vmf, vmjc, vmd, vmacc) = infer(args, model, val_queue, criterion) logger.info( "ValLoss:{:.3f} ValAcc:{:.3f} ValDice:{:.3f} ValJc:{:.3f}" .format(val_loss, vmacc, vmd, vmjc)) writer.add_scalar('Val/loss', val_loss, epoch) is_best = True if (vmjc >= max_value) else False max_value = max(max_value, vmjc) state = { 'epoch': epoch, 'optimizer_arch': optimizer_arch.state_dict(), 'optimizer_weight': optimizer_weight.state_dict(), 'scheduler': scheduler.state_dict(), 'state_dict': model.state_dict(), 'alphas_dict': model.alphas_dict(), } logger.info("epoch:{} best:{} max_value:{}".format( epoch, is_best, max_value)) if not is_best: torch.save( state, os.path.join(save_model_path, "checkpoint.pth.tar")) else: torch.save( state, os.path.join(save_model_path, "checkpoint.pth.tar")) torch.save( state, os.path.join(save_model_path, "model_best.pth.tar")) # one stage end, we should change the operations num (divided 2) weight_down = F.softmax(model.arch_parameters()[0], dim=-1).data.cpu().numpy() weight_up = F.softmax(model.arch_parameters()[1], dim=-1).data.cpu().numpy() weight_normal = F.softmax(model.arch_parameters()[2], dim=-1).data.cpu().numpy() weight_network = F.softmax(model.arch_parameters()[3], dim=-1).data.cpu().numpy() logger.info("alphas_down: \n{}".format(weight_down)) logger.info("alphas_up: \n{}".format(weight_up)) logger.info("alphas_normal: \n{}".format(weight_normal)) logger.info("alphas_network: \n{}".format(weight_network)) genotype = model.genotype() logger.info('Stage:{} \n Genotype: {}'.format(sp, genotype)) logger.info( '------Stage {} end ! Then Dropping Paths------'.format(sp)) # 6 4 7 # CellLinkDownPos CellLinkUpPos CellPos # # 6-->3-->1 # drop_op_down = [3, 2] # # 4-->2-->1 # drop_op_up = [2, 1] # # 7-->4-->1 # drop_op_normal = [3, 3] # update switches in 0 stage end if sp == 0: switches_down = update_switches(weight_down.copy(), switches_down.copy(), CellLinkDownPos, drop_op_down[sp]) switches_up = update_switches(weight_up.copy(), switches_up.copy(), CellLinkUpPos, drop_op_up[sp]) switches_normal = update_switches(weight_normal.copy(), switches_normal.copy(), CellPos, drop_op_normal[sp]) logger.info('switches_down = %s', switches_down) logger.info('switches_up = %s', switches_up) logger.info('switches_normal = %s', switches_normal) logging_switches(logger, switches_down, CellLinkDownPos) logging_switches(logger, switches_up, CellLinkUpPos) logging_switches(logger, switches_normal, CellPos) else: # sp==1 is the final stage, we don`t need the keep operations # because we has the model.genotype # show the final one op in 14 mixop switches_down = update_switches(weight_down.copy(), switches_down.copy(), CellLinkDownPos, drop_op_down[sp]) switches_up = update_switches(weight_up.copy(), switches_up.copy(), CellLinkUpPos, drop_op_up[sp]) switches_normal = update_switches_nozero(weight_normal.copy(), switches_normal.copy(), CellPos, drop_op_normal[sp]) logger.info('switches_down = %s', switches_down) logger.info('switches_up = %s', switches_up) logger.info('switches_normal = %s', switches_normal) logging_switches(logger, switches_down, CellLinkDownPos) logging_switches(logger, switches_up, CellLinkUpPos) logging_switches(logger, switches_normal, CellPos) writer.close()
def main(args): #################### init logger ################################### log_dir = './eval' + '/{}'.format(args.dataset) + '/{}'.format(args.model) logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Eval'.format(args.model)) # setting args.save_path = log_dir args.save_images = os.path.join(args.save_path, "images") if not os.path.exists(args.save_images): os.mkdir(args.save_images) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### val_loader = get_dataloder(args, split_flag="valid") ######################## init model ############################################ if args.model == "layer7_double_deep_ep1600_8lr4e-3": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'layer7_double_deep' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=9, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/cvc/layer7_double_deep_ep1600_8lr4e-3/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif args.model == "alpha0_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_stage1_double_deep_ep200' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/cvc/alpha0_8lr4e-3/model_best.pth.tar' state_dict = torch.load(args.model_path, map_location='cpu')['state_dict'] state_dict = remove_module(state_dict) model.load_state_dict(state_dict) elif args.model == "alpha0_5_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/cvc/alpha0_5_8lr4e-3/model_best.pth.tar' state_dict = torch.load(args.model_path, map_location='cpu')['state_dict'] state_dict = remove_module(state_dict) model.load_state_dict(state_dict) elif args.model == "alpha1_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/cvc/alpha1_8lr4e-3/model_best.pth.tar' state_dict = torch.load(args.model_path, map_location='cpu')['state_dict'] state_dict = remove_module(state_dict) model.load_state_dict(state_dict) else: raise NotImplementedError() setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info(genotype) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) infer(args, model, criterion, val_loader, logger, args.save_images)
def main(args): #################### init logger ################################### log_dir = './logs/' + '{}'.format(args.dataset) + '/{}_{}_{}'.format(args.model,args.note,time.strftime('%Y%m%d-%H%M%S')) logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Train'.format(args.model)) # setting args.save_path = log_dir args.save_tbx_log = args.save_path + '/tbx_log' writer = SummaryWriter(args.save_tbx_log) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### train_loader = get_dataloder(args, split_flag="train") val_loader = get_dataloder(args, split_flag="valid") ############init model ########################### if args.model == "layer7_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'layer7_double_deep' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=7, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "stage1_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_double_deep' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "stage1_nodouble_deep": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_deep' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "stage1_nodouble_deep_slim": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_deep' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPruneSlim( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "alpha1_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha1_stage1_double_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "alpha0_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_stage1_double_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) #isic trans elif args.model == "stage1_layer9_110epoch_double_deep_final": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) # just normaL cell keep elif args.model == "dd_normal": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPruneNormal( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) # normal+down elif args.model == "dd_normaldown": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPruneNormalDown( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) # normal+up elif args.model == "dd_normalup": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPruneNormalUp( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) # normal+up+down elif args.model == "alpha0_5_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) # abliation study of channel doubling and deepsupervision elif args.model == "alpha0_5_stage1_double_nodeep_ep80": args.deepsupervision = False args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_nodeep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "alpha0_5_stage1_nodouble_deep_ep80": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'alpha0_5_stage1_nodouble_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) elif args.model == "alpha0_5_stage1_nodouble_nodeep_ep80": args.deepsupervision = False args.double_down_channel = False args.genotype_name = 'alpha0_5_stage1_nodouble_nodeep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux ) if torch.cuda.device_count() > 1 and args.use_cuda: logger.info('use: %d gpus', torch.cuda.device_count()) model = nn.DataParallel(model) setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info(genotype) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) # init optimizer optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) # init schedulers Steplr scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epoch) # scheduler=torch.optim.lr_scheduler.StepLR(optimizer=optimizer,step_size=30,gamma=0.1,last_epoch=-1) ############################### check resume ######################### start_epoch = 0 if args.resume is not None: if os.path.isfile(args.resume): logger.info("Loading model and optimizer from checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume, map_location=args.device) start_epoch = checkpoint['epoch'] optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['state_dict']) scheduler.load_state_dict(checkpoint['scheduler']) else: raise FileNotFoundError("No checkpoint found at '{}'".format(args.resume)) #################################### train and val ######################## max_value = 0 for epoch in range(start_epoch, args.epoch): # lr=adjust_learning_rate(args,optimizer,epoch) scheduler.step() logger.info('Epoch: %d lr %e', epoch, scheduler.get_lr()[0]) # train if args.deepsupervision: mean_loss, value1, value2 = train(args, model, criterion, train_loader, optimizer) mr, ms, mp, mf, mjc, md, macc = value1 mmr, mms, mmp, mmf, mmjc, mmd, mmacc = value2 logger.info( "Epoch:{} Train_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}".format(epoch, mean_loss, macc, md, mjc)) logger.info(" dmAcc:{:.3f} dmDice:{:.3f} dmJc:{:.3f}".format(mmacc, mmd, mmjc)) writer.add_scalar('Train/dmAcc', mmacc, epoch) writer.add_scalar('Train/dRecall', mmr, epoch) writer.add_scalar('Train/dSpecifi', mms, epoch) writer.add_scalar('Train/dPrecision', mmp, epoch) writer.add_scalar('Train/dF1', mmf, epoch) writer.add_scalar('Train/dJc', mmjc, epoch) writer.add_scalar('Train/dDice', mmd, epoch) else: mean_loss, value1 = train(args, model, criterion, train_loader, optimizer) mr, ms, mp, mf, mjc, md, macc = value1 logger.info( "Epoch:{} Train_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}".format(epoch, mean_loss, macc, md, mjc)) # write writer.add_scalar('Train/Loss', mean_loss, epoch) writer.add_scalar('Train/mAcc', macc, epoch) writer.add_scalar('Train/Recall', mr, epoch) writer.add_scalar('Train/Specifi', ms, epoch) writer.add_scalar('Train/Precision', mp, epoch) writer.add_scalar('Train/F1', mf, epoch) writer.add_scalar('Train/Jc', mjc, epoch) writer.add_scalar('Train/Dice', md, epoch) # val if args.deepsupervision: vmean_loss, valuev1, valuev2 = infer(args, model, criterion, val_loader) vmr, vms, vmp, vmf, vmjc, vmd, vmacc = valuev1 mvmr, mvms, mvmp, mvmf, mvmjc, mvmd, mvmacc = valuev2 logger.info( "Epoch:{} Val_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}".format(epoch, vmean_loss, vmacc, vmd, vmjc)) logger.info(" dmAcc:{:.3f} dmDice:{:.3f} dmJc:{:.3f}".format(mvmacc, mvmd, mvmjc)) writer.add_scalar('Val/mAcc', mvmacc, epoch) writer.add_scalar('Val/Recall', mvmr, epoch) writer.add_scalar('Val/Specifi', mvms, epoch) writer.add_scalar('Val/Precision', mvmp, epoch) writer.add_scalar('Val/F1', mvmf, epoch) writer.add_scalar('Val/Jc', mvmjc, epoch) writer.add_scalar('Val/Dice', mvmd, epoch) else: vmean_loss, valuev1 = infer(args, model, criterion, val_loader) vmr, vms, vmp, vmf, vmjc, vmd, vmacc = valuev1 logger.info( "Epoch:{} Val_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}".format(epoch, vmean_loss, vmacc, vmd, vmjc)) is_best = True if (vmjc >=max_value) else False max_value = max(max_value, vmjc) writer.add_scalar('Val/Loss', vmean_loss, epoch) writer.add_scalar('Val/mAcc', vmacc, epoch) writer.add_scalar('Val/Recall', vmr, epoch) writer.add_scalar('Val/Specifi', vms, epoch) writer.add_scalar('Val/Precision', vmp, epoch) writer.add_scalar('Val/F1', vmf, epoch) writer.add_scalar('Val/Jc', vmjc, epoch) writer.add_scalar('Val/Dice', vmd, epoch) state={ 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'state_dict': model.state_dict(), 'scheduler': model.state_dict(), } logger.info("epoch:{} best:{} max_value:{}".format(epoch,is_best,max_value)) if not is_best: torch.save(state,os.path.join(args.save_path,"checkpoint.pth.tar")) else: torch.save(state,os.path.join(args.save_path,"checkpoint.pth.tar")) torch.save(state,os.path.join(args.save_path,"model_best.pth.tar")) writer.close()
def main(args): #################### init logger ################################### # args.model_list=["unet","unet++_deep","unet++_nodeep",'attention_unet_v1','multires_unet','r2unet_t3', # 'unet_ep800dice','unet++_deep_ep800dice','unet++_nodeep_ep800dice','attention_unet_v1_ep800dice','multires_unet_ep800dice' # ] args.model_list = ['unet', 'unet++', "attention_unet", "multires_unet"] for model_name in args.model_list: if model_name == 'unet': args.model = 'unet' model_weight_path = './logs/chaos/unet_ep150_v2/20200403-134703/checkpoint.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'unet++': args.model = 'unet++' args.deepsupervision = False model_weight_path = './logs/chaos/unet++_ep150_v2/20200403-135401/checkpoint.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) # elif model_name == 'unet++_deep': # args.model = 'unet++' # args.deepsupervision = True # model_weight_path = './logs/unet++_deep_ep1600/cvc/20200312-143345/model_best.pth.tar' # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'attention_unet': args.model = 'attention_unet_v1' args.deepsupervision = False model_weight_path = './logs/chaos/attention_unet_v1_ep150_v2/20200403-135445/checkpoint.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'multires_unet': args.model = 'multires_unet' args.deepsupervision = False model_weight_path = './logs/chaos/multires_unet_ep150_v2/20200403-135549/checkpoint.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) else: raise NotImplementedError() assert os.path.exists(args.save) args.model_save_path = os.path.join(args.save, model_name) logger = get_logger(args.model_save_path) args.save_images = os.path.join(args.model_save_path, "images") if not os.path.exists(args.save_images): os.mkdir(args.save_images) if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True val_loader = get_dataloder(args, split_flag="valid") setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) infer(args, model, criterion, val_loader, logger, args.save_images)
def main(args): #################### init logger ################################### args.model_list=["unet","unet++",'attention_unet_v1','multires_unet','r2unet_t3'] for model_name in args.model_list: if model_name=='unet': args.model='unet' model_weight_path='./logs/unet_ep1600/cvc/20200312-143050/model_best.pth.tar' model=get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name=='unet++': args.model='unet++' args.deepsupervision=False model_weight_path='./logs/unet++_ep1600/cvc/20200312-143358/model_best.pth.tar' model=get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'attention_unet_v1': args.model = 'attention_unet_v1' model_weight_path = './logs/attention_unet_v1_ep1600/cvc/20200312-143413/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'multires_unet': args.model = 'multires_unet' model_weight_path = './logs/multires_unet_ep1600_t2/20200322-194117/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # change bn relu order elif model_name == 'multires_unet_align': args.model = 'multires_unet' model_weight_path = './logs/multires_unet_ep1600_chbnrelu/20200327-184457/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'r2unet_t3': args.model = 'r2unet' args.time_step=3 model_weight_path = './logs/r2unet_ep1600_t2/20200324-032815/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'unet_ep800dice': args.model = 'unet' model_weight_path = './logs/unet_ep800_bcedice/cvc/20200315-043021/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name=='unet++_nodeep_ep800dice': args.model='unet++' args.deepsupervision=False model_weight_path='./logs/unet++_ep800_bcedice/cvc/20200315-043214/model_best.pth.tar' model=get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'unet++_deep_ep800dice': args.model = 'unet++' args.deepsupervision = True model_weight_path = './logs/unet++_deep_ep800_bcedice/cvc/20200315-043134/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'attention_unet_v1_ep800dice': args.model = 'attention_unet_v1' args.deepsupervision=False model_weight_path = './logs/attention_unet_v1_ep800_bcedice/cvc/20200315-043300/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'multires_unet_ep800dice': args.model = 'multires_unet' args.deepsupervision=False model_weight_path = './logs/multires_unet_ep800_bcedice/cvc/20200312-173031/model_best.pth.tar' model = get_models(args) model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) else: raise NotImplementedError() assert os.path.exists(args.save) args.model_save_path=os.path.join(args.save,model_name) logger = get_logger(args.model_save_path) args.save_images= os.path.join(args.model_save_path,"images") if not os.path.exists(args.save_images): os.mkdir(args.save_images) if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True val_loader = get_dataloder(args, split_flag="valid") setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) infer(args, model, criterion, val_loader,logger,args.save_images)
def main(args): #################### init logger ################################### #args.model_list=["unet","unet++_deep",'attention_unet_v1','multires_unet', 'r2unet_t3'] args.model_list = [ "unet", "unet++_deep", 'unet++_nodeep', "attention_unet_v1", "multires_unet", "r2unet" ] for model_name in args.model_list: # if model_name=='unet': # args.model='unet' # model_weight_path='./logs/isic/logs_coslr/unet/isic2018/20200229-035150/checkpoint.pth.tar' # model=get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # elif model_name=='unet++_deep': # args.model='unet++' # args.deepsupervision=True # model_weight_path='./logs/isic/logs_coslr/unet++/isic2018/20200229-035514/checkpoint.pth.tar' # model=get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # elif model_name == 'unet++_nodeep': # args.model = 'unet++' # args.deepsupervision = False # model_weight_path = '/checkpoint.pth.tar' # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # elif model_name == 'attention_unet_v1': # args.model = 'attention_unet_v1' # model_weight_path = './logs/isic/logs_coslr/attention_unet_v1/isic2018/20200302-190718/checkpoint.pth.tar' # args.deepsupervision=False # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # # elif model_name == 'multires_unet': # args.model = 'multires_unet' # model_weight_path = './logs/isic/logs_coslr/multires_unet/isic2018/20200229-035734/checkpoint.pth.tar' # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # # elif model_name == 'r2unet_t3': # args.model = 'r2unet' # args.time_step=3 # model_weight_path = './logs/isic/logs_coslr/r2unet/isic2018/20200302-190808/checkpoint.pth.tar' # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # ep300 baseline if model_name == 'unet': args.model = 'unet' model_weight_path = './logs/isic2018/unet_ep300/20200402-135108/model_best.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'unet++_deep': args.model = 'unet++' args.deepsupervision = True model_weight_path = './logs/isic2018/unet++_ep300_deep/20200402-135243/model_best.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'unet++_nodeep': args.model = 'unet++' args.deepsupervision = False model_weight_path = './logs/isic2018/unet++_ep300/20200402-135317/model_best.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'attention_unet_v1': args.model = 'attention_unet_v1' args.deepsupervision = False model_weight_path = './logs/isic2018/attention_unet_v1_ep300/20200413-160808//model_best.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'multires_unet': args.model = 'multires_unet' args.deepsupervision = False model_weight_path = './logs/isic2018/attention_unet_v1_ep300/20200413-160808//model_best.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) elif model_name == 'r2unet': args.model = 'r2unet' args.deepsupervision = False model_weight_path = './logs/isic2018/attention_unet_v1_ep300/20200413-160808//model_best.pth.tar' model = get_models(args) model.load_state_dict( torch.load(model_weight_path, map_location='cpu')['state_dict']) # elif model_name == 'attention_unet_v1': # args.model = 'attention_unet_v1' # model_weight_path = './logs/isic/logs_coslr/attention_unet_v1/isic2018/20200302-190718/checkpoint.pth.tar' # args.deepsupervision=False # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # # elif model_name == 'multires_unet': # args.model = 'multires_unet' # model_weight_path = './logs/isic/logs_coslr/multires_unet/isic2018/20200229-035734/checkpoint.pth.tar' # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) # # elif model_name == 'r2unet_t3': # args.model = 'r2unet' # args.time_step=3 # model_weight_path = './logs/isic/logs_coslr/r2unet/isic2018/20200302-190808/checkpoint.pth.tar' # model = get_models(args) # model.load_state_dict(torch.load(model_weight_path, map_location='cpu')['state_dict']) else: raise NotImplementedError() assert os.path.exists(args.save) args.model_save_path = os.path.join(args.save, model_name) logger = get_logger(args.model_save_path) args.save_images = os.path.join(args.model_save_path, "images") if not os.path.exists(args.save_images): os.mkdir(args.save_images) if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True val_loader = get_dataloder(args, split_flag="valid") setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) infer(args, model, criterion, val_loader, logger, args.save_images)
def main(args): args.model_list = [ 'double_deep', 'double', 'nodouble', 'nodouble_deep', 'slim_dd', 'slim_double', 'slim_nodouble', 'slim_nodouble_deep' ] #args.model_list=["slim_nodouble_deep_init32"] for model_name in args.model_list: print(model_name) if model_name == "double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/prune_20200313-063406_32_32_ep300_double_deep/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'double': args.deepsupervision = False args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/prune_20200313-063428_32_32_ep300_double/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'nodouble': args.deepsupervision = False args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/prune_20200316-141125_nodouble_32_ep300/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'nodouble_deep': args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/prune_20200316-141242_nodouble_32_ep300_deep/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) if model_name == "slim_dd": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = net_dd(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/dd_20200319-170442_ep300/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'slim_double': args.deepsupervision = False args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_final' genotype = eval('genotypes.%s' % args.genotype_name) model = net_double(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/double_20200319-170621_ep300/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'slim_nodouble': args.deepsupervision = False args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_final' genotype = eval('genotypes.%s' % args.genotype_name) model = net_nodouble(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/nodouble_20200319-210910_ep300/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'slim_nodouble_deep': args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = net_nodouble_deep( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/nodouble_deep_20200319-210600_ep300/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == 'slim_nodouble_deep_init32': args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = net_nodouble_deep( genotype=genotype, input_c=args.in_channels, c=32, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/nodouble_deep_ep300_init32/model_best.pth.tar' model.load_state_dict( torch.load(args.model_path, map_location='cpu')['state_dict']) #################### init logger ################################### log_dir = './eval' + '/{}'.format( args.dataset) + '/{}'.format(model_name) ##################### init model ######################################## logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Eval'.format(model_name)) # setting args.save_path = log_dir args.save_images = os.path.join(args.save_path, "images") if not os.path.exists(args.save_images): os.mkdir(args.save_images) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### # sorted vaild datasets val_loader = get_dataloder(args, split_flag="valid") setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info(genotype) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device)
def main(args): #################### init logger ################################### log_dir = './logs/' + '{}'.format(args.dataset) + '/{}_{}_{}'.format( args.model, args.note, time.strftime('%Y%m%d-%H%M%S')) logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Train'.format(args.model)) # setting setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) args.save_path = log_dir args.save_tbx_log = args.save_path + '/tbx_log' writer = SummaryWriter(args.save_tbx_log) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### train_loader = get_dataloder(args, split_flag="train") val_loader = get_dataloder(args, split_flag="valid") ######################## init model ############################################ # model ############init model ########################### if args.model == "nodouble_deep_init32_ep100": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'nodouble_deep_init32_ep100' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=32, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=9, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "nodouble_deep_isic": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "nodouble_deep_drop02_layer7end": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'nodouble_deep_drop02_layer7end' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "stage1_nodouble_deep_ep36": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_nodouble_deep_ep36' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "stage1_nodouble_deep_ep63": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_nodouble_deep_ep63' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "stage1_nodouble_deep_ep83": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_nodouble_deep_ep83' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha1_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha1_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_5_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) # isic trans elif args.model == "stage1_layer9_110epoch_double_deep_final": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_deep_final' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) #chaos elif args.model == "stage0_double_deep_ep80_newim": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage0_double_deep_ep80_newim' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "stage1_double_deep_ep80_ts": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_double_deep_ep80_ts' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) # cvc trans elif args.model == "layer7_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'layer7_double_deep' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) if torch.cuda.device_count() > 1 and args.use_cuda: logger.info('use: %d gpus', torch.cuda.device_count()) model = nn.DataParallel(model) setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info(genotype) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() elif args.loss == 'multibcedice': criterion = MultiClassEntropyDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) # init optimizer optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) # init schedulers Steplr scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, args.epoch) # scheduler=torch.optim.lr_scheduler.StepLR(optimizer=optimizer,step_size=30,gamma=0.1,last_epoch=-1) ############################### check resume ######################### start_epoch = 0 if args.resume is not None: if os.path.isfile(args.resume): logger.info( "Loading model and optimizer from checkpoint '{}'".format( args.resume)) checkpoint = torch.load(args.resume, map_location=args.device) start_epoch = checkpoint['epoch'] optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['state_dict']) scheduler.load_state_dict(checkpoint['scheduler']) else: raise FileNotFoundError("No checkpoint found at '{}'".format( args.resume)) #################################### train and val ######################## max_value = 0 for epoch in range(start_epoch, args.epoch): # lr=adjust_learning_rate(args,optimizer,epoch) scheduler.step() # logger.info('Epoch: %d lr %e', epoch, scheduler.get_lr()[0]) # train total_loss = train(args, model, criterion, train_loader, optimizer, epoch, logger) # write writer.add_scalar('Train/total_loss', total_loss, epoch) # val tloss, md = val(args, model, criterion, val_loader, epoch, logger) writer.add_scalar('Val/total_loss', tloss, epoch) is_best = True if (md >= max_value) else False max_value = max(max_value, md) state = { 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'state_dict': model.state_dict(), 'scheduler': model.state_dict(), } logger.info("epoch:{} best:{} max_value:{}".format( epoch, is_best, max_value)) if not is_best: torch.save(state, os.path.join(args.save_path, "checkpoint.pth.tar")) else: torch.save(state, os.path.join(args.save_path, "checkpoint.pth.tar")) torch.save(state, os.path.join(args.save_path, "model_best.pth.tar")) writer.close()
def main(args): #################### init logger ################################### log_dir = './logs/' + '{}'.format(args.dataset) + '/{}_{}_{}'.format( args.model, time.strftime('%Y%m%d-%H%M%S'), args.note) logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Train'.format(args.model)) # setting args.save_path = log_dir args.save_tbx_log = args.save_path + '/tbx_log' writer = SummaryWriter(args.save_tbx_log) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### train_loader = get_dataloder(args, split_flag="train") val_loader = get_dataloder(args, split_flag="valid") ######################## init model ############################################ # model # get the network parameters if args.model == "alpha_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_deep_final' args.alphas_model = './search_exp/Nas_Search_Unet/isic2018/deepsupervision/stage_1_model/checkpoint.pth.tar' model_alphas = torch.load( args.alphas_model, map_location=args.device)['alphas_dict']['alphas_network'] model_alphas.requires_grad = False model_alphas = F.softmax(model_alphas, dim=-1) genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnet(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha_double": args.deepsupervision = False args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_final' args.alphas_model = './search_exp/Nas_Search_Unet/isic2018/nodeepsupervision/stage_1_model/checkpoint.pth.tar' model_alphas = torch.load( args.alphas_model, map_location=args.device)['alphas_dict']['alphas_network'] model_alphas.requires_grad = False model_alphas = F.softmax(model_alphas, dim=-1) genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnet(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha_nodouble": args.deepsupervision = False args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_final' args.alphas_model = './search_exp/Nas_Search_Unet/isic2018/nodouble/stage_1_model/checkpoint.pth.tar' model_alphas = torch.load( args.alphas_model, map_location=args.device)['alphas_dict']['alphas_network'] model_alphas.requires_grad = False model_alphas = F.softmax(model_alphas, dim=-1) genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnet(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha_nodouble_deep": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_deep_final' args.alphas_model = './search_exp/Nas_Search_Unet/isic2018/nodouble_deep/stage_1_model/checkpoint.pth.tar' model_alphas = torch.load( args.alphas_model, map_location=args.device)['alphas_dict']['alphas_network'] model_alphas.requires_grad = False model_alphas = F.softmax(model_alphas, dim=-1) genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnet(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_deep_final' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "double": args.deepsupervision = False args.double_down_channel = True args.genotype_name = 'stage1_layer9_110epoch_double_final' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "nodouble": args.deepsupervision = False args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_final' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "nodouble_deep": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'stage1_layer9_110epoch_deep_final' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha1_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha1_stage1_double_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_stage1_double_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_5_stage1_double_deep_ep80": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_5_stage1_double_nodeep_ep80": args.deepsupervision = False args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_nodeep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_5_stage1_nodouble_deep_ep80": args.deepsupervision = True args.double_down_channel = False args.genotype_name = 'alpha0_5_stage1_nodouble_deep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) elif args.model == "alpha0_5_stage1_nodouble_nodeep_ep80": args.deepsupervision = False args.double_down_channel = False args.genotype_name = 'alpha0_5_stage1_nodouble_nodeep_ep80' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) # cvc trans elif args.model == "layer7_double_deep": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'layer7_double_deep' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) # chaos trans elif args.model == "stage0_double_deep_ep80_newim": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'stage0_double_deep_ep80_newim' model_alphas = None genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune(genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) if torch.cuda.device_count() > 1 and args.use_cuda: logger.info('use: %d gpus', torch.cuda.device_count()) model = nn.DataParallel(model) setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info(genotype) logger.info(model_alphas) flop, param = get_model_complexity_info(model, (3, 256, 256), as_strings=True, print_per_layer_stat=False) print("GFLOPs: {}".format(flop)) print("Params: {}".format(param)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) # init optimizer optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum) # init schedulers Steplr scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, args.epoch) #scheduler=torch.optim.lr_scheduler.StepLR(optimizer=optimizer,step_size=30,gamma=0.1,last_epoch=-1) ############################### check resume ######################### start_epoch = 0 if args.resume is not None: if os.path.isfile(args.resume): logger.info( "Loading model and optimizer from checkpoint '{}'".format( args.resume)) checkpoint = torch.load(args.resume, map_location=args.device) start_epoch = checkpoint['epoch'] optimizer.load_state_dict(checkpoint['optimizer']) model.load_state_dict(checkpoint['state_dict']) scheduler.load_state_dict(checkpoint['scheduler']) else: raise FileNotFoundError("No checkpoint found at '{}'".format( args.resume)) #################################### train and val ######################## max_value = 0 for epoch in range(start_epoch, args.epoch): # lr=adjust_learning_rate(args,optimizer,epoch) scheduler.step() # train if args.deepsupervision: mean_loss, value1, value2 = train(args, model_alphas, model, criterion, train_loader, optimizer) mr, ms, mp, mf, mjc, md, macc = value1 logger.info( "Epoch:{} Train_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}". format(epoch, mean_loss, macc, md, mjc)) writer.add_scalar('Train/dDice', mmd, epoch) else: mean_loss, value1 = train(args, model_alphas, model, criterion, train_loader, optimizer) mr, ms, mp, mf, mjc, md, macc = value1 logger.info( "Epoch:{} Train_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}". format(epoch, mean_loss, macc, md, mjc)) # write writer.add_scalar('Train/Loss', mean_loss, epoch) # val if args.deepsupervision: vmean_loss, valuev1, valuev2 = infer(args, model_alphas, model, criterion, val_loader) vmr, vms, vmp, vmf, vmjc, vmd, vmacc = valuev1 logger.info( "Epoch:{} Val_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}". format(epoch, vmean_loss, vmacc, vmd, vmjc)) else: vmean_loss, valuev1 = infer(args, model_alphas, model, criterion, val_loader) vmr, vms, vmp, vmf, vmjc, vmd, vmacc = valuev1 logger.info( "Epoch:{} Val_Loss:{:.3f} Acc:{:.3f} Dice:{:.3f} Jc:{:.3f}". format(epoch, vmean_loss, vmacc, vmd, vmjc)) is_best = True if vmjc >= max_value else False max_value = max(max_value, vmjc) writer.add_scalar('Val/Loss', vmean_loss, epoch) state = { 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'state_dict': model.state_dict(), 'scheduler': model.state_dict(), } logger.info("epoch:{} best:{} max_value:{}".format( epoch, is_best, max_value)) if not is_best: torch.save(state, os.path.join(args.save_path, "checkpoint.pth.tar")) else: torch.save(state, os.path.join(args.save_path, "checkpoint.pth.tar")) torch.save(state, os.path.join(args.save_path, "model_best.pth.tar")) writer.close()
def main(args): #args.model_list=['alpha0_double_deep_0.01','alpha0_5_double_deep_0.01','alpha1_double_deep_0.01','nodouble_deep','slim_dd','slim_double','slim_nodouble','slim_nodouble_deep'] #args.model_list=["double_deep","nodouble_deep","slim_nodouble"] #args.model_list=["slim_nodouble_deep_init32"] #args.model_list=["slim_nodouble_deep_init48"] args.model_list = [ 'alpha0_double_deep_0.01', 'alpha0_5_double_deep_0.01', 'alpha1_double_deep_0.01' ] for model_name in args.model_list: if model_name == "alpha0_double_deep_0.01": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_stage1_double_deep_ep200' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/alpha0_double_deep_0.01/model_best.pth.tar' # kwargs = {'map_location': lambda storage, loc: storage.cuda(0)} # state_dict = torch.load(args.model_path, **kwargs) # # create new OrderedDict that does not contain `module.` # model.load_state_dict(state_dict) state_dict = torch.load(args.model_path, map_location='cpu')['state_dict'] state_dict = remove_module(state_dict) model.load_state_dict(state_dict) elif model_name == "alpha0_5_double_deep_0.01": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha0_5_stage1_double_deep_ep80' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/alpha0_5_double_deep_0.01/model_best.pth.tar' state_dict = torch.load(args.model_path, map_location='cpu')['state_dict'] state_dict = remove_module(state_dict) model.load_state_dict(state_dict) #model.load_state_dict(torch.load(args.model_path, map_location='cpu')['state_dict']) elif model_name == "alpha1_double_deep_0.01": args.deepsupervision = True args.double_down_channel = True args.genotype_name = 'alpha1_stage1_double_deep_ep200' genotype = eval('genotypes.%s' % args.genotype_name) model = BuildNasUnetPrune( genotype=genotype, input_c=args.in_channels, c=args.init_channels, num_classes=args.nclass, meta_node_num=args.middle_nodes, layers=args.layers, dp=args.dropout_prob, use_sharing=args.use_sharing, double_down_channel=args.double_down_channel, aux=args.aux) args.model_path = './logs/isic2018/alpha1_double_deep_0.01/model_best.pth.tar' state_dict = torch.load(args.model_path, map_location='cpu')['state_dict'] state_dict = remove_module(state_dict) model.load_state_dict(state_dict) #model.load_state_dict(torch.load(args.model_path, map_location='cpu')['state_dict']) #################### init logger ################################### log_dir = './eval' + '/{}'.format( args.dataset) + '/{}'.format(model_name) ##################### init model ######################################## logger = get_logger(log_dir) print('RUNDIR: {}'.format(log_dir)) logger.info('{}-Eval'.format(model_name)) # setting args.save_path = log_dir args.save_images = os.path.join(args.save_path, "images") if not os.path.exists(args.save_images): os.mkdir(args.save_images) ##################### init device ################################# if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) np.random.seed(args.manualSeed) torch.manual_seed(args.manualSeed) args.use_cuda = args.gpus > 0 and torch.cuda.is_available() args.device = torch.device('cuda' if args.use_cuda else 'cpu') if args.use_cuda: torch.cuda.manual_seed(args.manualSeed) cudnn.benchmark = True ####################### init dataset ########################################### # sorted vaild datasets val_loader = get_dataloder(args, split_flag="valid") setting = {k: v for k, v in args._get_kwargs()} logger.info(setting) logger.info(genotype) logger.info('param size = %fMB', calc_parameters_count(model)) # init loss if args.loss == 'bce': criterion = nn.BCELoss() elif args.loss == 'bcelog': criterion = nn.BCEWithLogitsLoss() elif args.loss == "dice": criterion = DiceLoss() elif args.loss == "softdice": criterion = SoftDiceLoss() elif args.loss == 'bcedice': criterion = BCEDiceLoss() else: criterion = nn.CrossEntropyLoss() if args.use_cuda: logger.info("load model and criterion to gpu !") model = model.to(args.device) criterion = criterion.to(args.device) infer(args, model, criterion, val_loader, logger, args.save_images)