all_dev = parse_devices(args.devices) network = PSPNet(config.num_classes) data_setting = { 'img_root': config.img_root_folder, 'gt_root': config.gt_root_folder, 'train_source': config.train_source, 'eval_source': config.eval_source, 'test_source': config.test_source } dataset = Cil(data_setting, 'val', None) if args.speed_test: device = all_dev[0] logger.info("=========DEVICE:%s SIZE:%s=========" % (torch.cuda.get_device_name(device), args.input_size)) input_size = tuple(int(x) for x in args.input_size.split('x')) compute_speed(network, input_size, device, args.iteration) elif args.summary: input_size = tuple(int(x) for x in args.input_size.split('x')) stat(network, input_size) else: with torch.no_grad(): segmentor = SegEvaluator(dataset, config.num_classes, config.image_mean, config.image_std, network, config.eval_scale_array, config.eval_flip, all_dev, args.verbose, args.save_path) segmentor.run(config.snapshot_dir, args.epochs, config.val_log_file, config.link_val_log_file)
def main(): global args, best_prec1, USE_GPU, device args = parser.parse_args() with open(args.config) as f: config = yaml.load(f) for k, v in config['common'].items(): setattr(args, k, v) ## Set random seeds ## torch.manual_seed(args.random_seed) np.random.seed(args.random_seed) random.seed(args.random_seed) #print(args) #print(xxx) # create models if args.input_size != 224 or args.image_size != 256: image_size = args.image_size input_size = args.input_size else: image_size = 256 input_size = 224 print("Input image size: {}, test size: {}".format(image_size, input_size)) if "model" in config.keys(): model = models.__dict__[args.arch](**config['model']) else: model = models.__dict__[args.arch]() device = torch.device( 'cuda:' + str(args.gpus[0]) if torch.cuda.is_available() else "cpu") str_input_size = '1x3x224x224' if args.summary: input_size = tuple(int(x) for x in str_input_size.split('x')) stat(model, input_size) return if USE_GPU: cudnn.benchmark = True torch.cuda.manual_seed_all(args.random_seed) args.gpus = [int(i) for i in args.gpus.split(',')] model = torch.nn.DataParallel(model, device_ids=args.gpus) model.to(device) count_params(model) pytorch_total_params = sum(p.numel() for p in model.parameters()) print('total_params', pytorch_total_params) # define loss function (criterion) and optimizer criterion = FocalLoss(device, 2, gamma=args.fl_gamma) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if args.speed: input_size = tuple(int(x) for x in str_input_size.split('x')) iteration = 1000 compute_speed(model, input_size, device, iteration) return ########################################################## # adjust start learning rate args.lr = args.lr * (args.batch_size // 32) ######################################################### # optionally resume from a checkpoint if args.resume: print(os.getcwd()) if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) # Data loading code normalize = transforms.Normalize( mean=[0.14300402, 0.1434545, 0.14277956], ##accorcoding to casia-surf val to commpute std=[0.10050353, 0.100842826, 0.10034215]) img_size = args.input_size ratio = 224.0 / float(img_size) train_dataset = CASIA( args, transforms.Compose([ transforms.RandomResizedCrop(img_size), transforms.RandomHorizontalFlip(), #ColorTransform(p=0.3), transforms.ToTensor(), ColorAugmentation(), normalize, ]), phase_train=True) val_dataset = CASIA(args, transforms.Compose([ transforms.Resize(int(256 * ratio)), transforms.CenterCrop(img_size), transforms.ToTensor(), normalize, ]), phase_train=False, phase_test=args.phase_test) train_sampler = None val_sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=(train_sampler is None), sampler=train_sampler) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler) if args.evaluate: validate(val_loader, model, criterion, args.start_epoch) return else: print(model) for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) # train for one epoch train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set prec1 = validate(val_loader, model, criterion, epoch) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 if is_best: print('epoch: {} The best is {} last best is {}'.format( epoch, prec1, best_prec1)) best_prec1 = max(prec1, best_prec1) if not os.path.exists(args.save_path): os.makedirs(args.save_path) if is_best: save_name = '{}/{}_{}_best.pth.tar'.format(args.save_path, args.model_name, epoch) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, filename=save_name)
def main(): global args, best_prec1, USE_GPU, device, sub, modal args = parser.parse_args() modal = args.modal SUB = ['4@1', '4@2', '4@3'] sub = SUB[args.sub - 1] with open(args.config) as f: config = yaml.load(f) for k, v in config['common'].items(): setattr(args, k, v) ## Set random seeds ## torch.manual_seed(args.random_seed) np.random.seed(args.random_seed) random.seed(args.random_seed) # create models if args.input_size != 224 or args.image_size != 256: image_size = args.image_size input_size = args.input_size else: image_size = 256 input_size = 224 print("Input image size: {}, test size: {}".format(image_size, input_size)) if "model" in config.keys(): model = models.__dict__[args.arch](**config['model']) else: model = models.__dict__[args.arch]() device = torch.device('cuda:' + str(args.gpus[0]) if torch.cuda.is_available() else "cpu") str_input_size = '1x3x224x224' if args.summary: input_size = tuple(int(x) for x in str_input_size.split('x')) stat(model, input_size) return if USE_GPU: cudnn.benchmark = True torch.cuda.manual_seed_all(args.random_seed) args.gpus = [int(i) for i in args.gpus.split(',')] model = torch.nn.DataParallel(model, device_ids=args.gpus) model.to(device) # count_params(model) pytorch_total_params = sum(p.numel() for p in model.parameters()) print('total_params', pytorch_total_params) # define loss function (criterion) and optimizer criterion = FocalLoss(device, 2, gamma=args.fl_gamma) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if args.speed: input_size = tuple(int(x) for x in str_input_size.split('x')) iteration = 1000 compute_speed(model, input_size, device, iteration) return # optionally resume from a checkpoint if args.resume: print(os.getcwd()) if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) norm = dict() norm["4@1_depth"] = transforms.Normalize(mean=[0.70745318, 0.70745318, 0.70745318], std=[0.26528493, 0.26528493, 0.26528493]) norm["4@2_depth"] = transforms.Normalize(mean=[0.67922, 0.67922, 0.67922], std=[0.2890028, 0.2890028, 0.2890028]) norm["4@3_depth"] = transforms.Normalize(mean=[0.7110071, 0.7110071, 0.7110071], std=[0.28458922, 0.28458922, 0.28458922]) norm["4@1_ir"] = transforms.Normalize(mean=[0.22784027, 0.22784027, 0.22784027], std=[0.10182471, 0.10182471, 0.10182471]) norm["4@2_ir"] = transforms.Normalize(mean=[0.34846747, 0.34846747, 0.34846747], std=[0.17954783, 0.17954783, 0.17954783]) norm["4@3_ir"] = transforms.Normalize(mean=[0.31522677, 0.31522677, 0.31522677], std=[0.16015441, 0.16015441, 0.16015441]) img_size = args.input_size if modal == 'merge': norm_type = ['ir', 'depth'] else: norm_type = [modal] train_trans = dict() val_trans = dict() for n in norm_type: train_trans[n] = transforms.Compose([ transforms.RandomResizedCrop(img_size), transforms.RandomHorizontalFlip(), transforms.RandomRotation(10), transforms.ColorJitter(0.1, 0.1), transforms.ToTensor(), # ColorAugmentation(), norm["{}_{}".format(sub, n)], ]) val_trans[n] = transforms.Compose([ transforms.Resize([224, 224]), transforms.ToTensor(), norm["{}_{}".format(sub, n)], ]) train_dataset = CASIA( train_trans, phase_train=True, sub=sub, modal=modal) val_dataset = CASIA( val_trans, phase_train=False, phase_test=args.phase_test, sub=sub, modal=modal) train_sampler = None val_sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=(train_sampler is None), sampler=train_sampler) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False, sampler=val_sampler) if args.evaluate: validate(val_loader, model, criterion, args.start_epoch) return else: print(model) for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) # train for one epoch train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set if args.train_val: prec1 = validate(val_loader, model, criterion, epoch) # remember best prec@1 and save checkpoints is_best = prec1 > best_prec1 if is_best: print('epoch: {} The best is {} last best is {}'.format(epoch, prec1, best_prec1)) best_prec1 = max(prec1, best_prec1) if not os.path.exists(args.save_path): os.mkdir(args.save_path) save_name = '{}/{}_{}_{}_{}_best.pth.tar'.format(args.save_path, sub, modal, args.model_name, epoch) if is_best else \ '{}/{}_{}_{}_{}.pth.tar'.format(args.save_path, sub, modal, args.model_name, epoch) save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, filename=save_name) else: if not os.path.exists(args.save_path): os.mkdir(args.save_path) save_name = '{}/{}_{}_{}_{}.pth.tar'.format(args.save_path, sub, modal, args.model_name, epoch) save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, filename=save_name)
def main(): global args, best_prec1, USE_GPU, device ## Set random seeds ## torch.manual_seed(args.random_seed) np.random.seed(args.random_seed) random.seed(args.random_seed) # create models if args.input_size != 224 or args.image_size != 256: image_size = args.image_size input_size = args.input_size else: image_size = 114 input_size = 112 logger.Print("Input image size: {}, test size: {}".format( image_size, input_size)) logger.Print(f"lr = {args.lr} batch_size = {args.batch_size}") logger.Print(f'protoal:{args.protoal}') model = Model(pretrained=False) device = torch.device( 'cuda:' + str(args.gpus[0]) if torch.cuda.is_available() else "cpu") str_input_size = '1x3x112x112' if args.summary: input_size = tuple(int(x) for x in str_input_size.split('x')) stat(model, input_size) return if USE_GPU: cudnn.benchmark = True torch.cuda.manual_seed_all(args.random_seed) args.gpus = [int(i) for i in args.gpus.split(',')] #model = torch.nn.DataParallel(model,device_ids=args.gpus) model.to(device) pytorch_total_params = sum(p.numel() for p in model.parameters()) logger.Print(f'total_params:{pytorch_total_params}') # define loss function (criterion) and optimizer if args.loss == "focal": criterion = FocalLoss(device, 2, gamma=args.fl_gamma) else: criterion = nn.CrossEntropyLoss() #criterion = criterion.to(device=config.device) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if args.speed: input_size = tuple(int(x) for x in str_input_size.split('x')) iteration = 1000 compute_speed(model, input_size, device, iteration) return # optionally resume from a checkpoint if args.resume: logger.Print(os.getcwd()) if os.path.isfile(args.resume): logger.Print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.Print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: logger.Print("=> no checkpoint found at '{}'".format(args.resume)) img_size = args.input_size train_loader, val_loader = load_cisia_csfa(root=args.train_path, protoal=args.protoal, img_size=img_size, train_batch=args.batch_size) if args.evaluate: validate(val_loader, model, criterion, args.start_epoch) return for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) # train for one epoch train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set prec1 = validate(val_loader, model, criterion, epoch) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 if is_best: logger.Print('epoch: {} The best is {} last best is {}'.format( epoch, prec1, best_prec1)) best_prec1 = max(prec1, best_prec1) save_path = log_path + '/models/' if not os.path.exists(save_path): os.mkdir(save_path) elif epoch % 2 == 0: save_name = '{}/{}_{}.pth.tar'.format(log_path, args.arch, epoch) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, filename=save_name)