def provide_model(model_name, classes_count, device, full=True): """ provides the pretrained model by model_name and replaces last classification layer with number of classes if full is provided and set to False, then the weights of the pretrained are set unmodifiable """ if model_name == 'hrnet': hrnet_model = load_hrnet_pretrained(device) num_ftrs = hrnet_model.classifier.in_features hrnet_model.classifier = nn.Linear(num_ftrs, classes_count) model = nn.Sequential(hrnet_model).to(device) elif model_name == 'resnet18': resnet18 = torchvision.models.resnet18(pretrained=True) num_ftrs = resnet18.fc.in_features resnet18.fc = nn.Linear(num_ftrs, classes_count) model = resnet18 elif model_name == 'resnet152': resnet152 = torchvision.models.resnet152(pretrained=True) num_ftrs = resnet152.fc.in_features resnet152.fc = nn.Linear(num_ftrs, classes_count) model = resnet152 dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) print(get_model_summary(model.cpu(), dump_input.cpu(), verbose=True)) return model
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'valid') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config) dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) logger.info(get_model_summary(model, dump_input)) if config.TEST.MODEL_FILE: logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE)) # model.load_state_dict(torch.load(config.TEST.MODEL_FILE)) model.load_state_dict({ k.replace('module.', ''): v for k, v in torch.load(config.TEST.MODEL_FILE) ['state_dict'].items() }) else: model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info('=> loading model from {}'.format(model_state_file)) model.load_state_dict(torch.load(model_state_file)) gpus = [0, 1] model = torch.nn.DataParallel(model, device_ids=gpus).cuda() # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() # Data loading code valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) test_dataset = datasets(dataset_root='./Data/', split='test', size=config.MODEL.IMAGE_SIZE[0]) valid_loader = torch.utils.data.DataLoader( test_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True) # evaluate on validation set validate(config, valid_loader, model, criterion, final_output_dir, tb_log_dir, None)
def main(): mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225] args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.' + config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: # model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth') model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load( '/home/hwits/Documents/CarVid/DDRNet/DDRNet.pytorch/model_best_bacc.pth.tar' ) if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] newstate_dict = { k: v for k, v in pretrained_dict.items() if k in model.state_dict() } # print(pretrained_dict.keys()) model.load_state_dict(newstate_dict) model = model.cuda() if True: save_wts = True print(model) if save_wts: f = open('DDRNetLite.wts', 'w') f.write('{}\n'.format(len(model.state_dict().keys()))) for k, v in model.state_dict().items(): print("Layer {} ; Size {}".format(k, v.cpu().numpy().shape)) vr = v.reshape(-1).cpu().numpy() f.write('{} {} '.format(k, len(vr))) for vv in vr: f.write(' ') f.write(struct.pack('>f', float(vv)).hex()) f.write('\n')
def validate_hrnet(): # cudnn related setting torch.backends.cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config) device = torch.device("cuda:0") dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) print( get_model_summary(model.cuda(device), dump_input.cuda(device), verbose=True)) if config.TEST.MODEL_FILE: print('=> loading model from {}'.format(config.TEST.MODEL_FILE)) model.load_state_dict(torch.load(config.TEST.MODEL_FILE)) else: model_state_file = os.path.join("./", 'final_state.pth.tar') print('=> loading model from {}'.format(model_state_file)) model.load_state_dict(torch.load(model_state_file)) # model = torch.nn.Sequential(model).cuda() model = torch.nn.Sequential(model).cuda(device) # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda(device) # Data loading code valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) dataset = datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ])) print(f'number of images: {len(dataset.imgs)}') valid_loader = torch.utils.data.DataLoader( dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU, shuffle=False, num_workers=config.WORKERS, pin_memory=True) # evaluate on validation set validate(config, valid_loader, model, criterion, './log', '/output', None)
def main(): mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] args = parse_args() logger, final_output_dir, _ = create_logger( config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.'+config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.'+config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]) ) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: # model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth') model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load('/home/kong/Documents/DDRNet.Pytorch/DDRNet.Pytorch/output/face/ddrnet23_slim/checkpoint.pth.tar') if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] newstate_dict = {k:v for k,v in pretrained_dict.items() if k in model.state_dict()} # print(pretrained_dict.keys()) model.load_state_dict(newstate_dict) model = model.to("cpu") print(model) model.eval() example = torch.rand(1, 3, 512, 512) model = torch.quantization.convert(model) # traced_script_module = torch.jit.trace(model, example) # traced_script_module.save("ddrnetfp32.pt") scriptedm = torch.jit.script(model) opt_model = torch.utils.optimize_for_mobile(scriptedm) torch.jit.save(opt_model, "ddrnetint8.pt")
def load_hrnet_pretrained(device): # cudnn related setting torch.backends.cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config) dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) if device.type == 'cpu': print(get_model_summary(model, dump_input, verbose=True)) else: print( get_model_summary(model.cuda(device), dump_input.cuda(device), verbose=True)) print('=> loading model from {}'.format(config.TEST.MODEL_FILE)) model.load_state_dict(torch.load(config.TEST.MODEL_FILE)) return model
def main(): args = parse_args() if args.dist_url == "env://" and args.world_size == -1: args.world_size = int(os.environ["WORLD_SIZE"]) dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url) logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config) dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) logger.info(get_model_summary(model, dump_input)) # copy model file # this_dir = os.path.dirname(__file__) # models_dst_dir = os.path.join(final_output_dir, 'models') # if os.path.exists(models_dst_dir): # shutil.rmtree(models_dst_dir) # shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) ''' model = torch.nn.DataParallel(model, device_ids=gpus).cuda() ''' # Change DP to DDP torch.cuda.set_device(args.local_rank) model = model.to(args.local_rank) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.local_rank], output_device=args.local_rank) # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() optimizer = get_optimizer(config, model) best_perf = 0.0 best_model = False last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_perf = checkpoint['perf'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) best_model = True if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) # Data loading code traindir = os.path.join(config.DATASET.ROOT, config.DATASET.TRAIN_SET) valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ''' train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) ) ''' # Change to TSV dataset instance train_dataset = TSVInstance( traindir, transforms.Compose([ transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) # DDP requires DistributedSampler train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=(train_sampler is None), num_workers=config.WORKERS, pin_memory=True, sampler=train_sampler) valid_loader = torch.utils.data.DataLoader( TSVInstance( valdir, transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ])), batch_size=config.TEST.BATCH_SIZE_PER_GPU, shuffle=False, num_workers=config.WORKERS, pin_memory=True) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() # train for one epoch train(config, train_loader, model, criterion, optimizer, epoch, final_output_dir, tb_log_dir, writer_dict) # evaluate on validation set perf_indicator = validate(config, valid_loader, model, criterion, final_output_dir, tb_log_dir, writer_dict) if perf_indicator > best_perf: best_perf = perf_indicator best_model = True else: best_model = False logger.info('=> saving checkpoint to {}'.format(final_output_dir)) save_checkpoint( { 'epoch': epoch + 1, 'model': config.MODEL.NAME, 'state_dict': model.module.state_dict(), 'perf': perf_indicator, 'optimizer': optimizer.state_dict(), }, best_model, final_output_dir, filename='checkpoint.pth.tar') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): args = parse_args() print(colored("Setting default tensor type to cuda.FloatTensor", "cyan")) torch.multiprocessing.set_start_method('spawn') torch.set_default_tensor_type('torch.cuda.FloatTensor') logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config).cuda() dump_input = torch.rand(config.TRAIN.BATCH_SIZE_PER_GPU, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]).cuda() logger.info(get_model_summary(model, dump_input)) if config.TRAIN.MODEL_FILE: model.load_state_dict(torch.load(config.TRAIN.MODEL_FILE)) logger.info( colored('=> loading model from {}'.format(config.TRAIN.MODEL_FILE), 'red')) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() print("Finished constructing model!") # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = get_optimizer(config, model) lr_scheduler = None best_perf = 0.0 best_model = False last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_perf = checkpoint['perf'] model.module.load_state_dict(checkpoint['state_dict']) # Update weight decay if needed checkpoint['optimizer']['param_groups'][0][ 'weight_decay'] = config.TRAIN.WD optimizer.load_state_dict(checkpoint['optimizer']) if 'lr_scheduler' in checkpoint: lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, 1e5, last_epoch=checkpoint['lr_scheduler']['last_epoch']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) best_model = True # Data loading code dataset_name = config.DATASET.DATASET if dataset_name == 'imagenet': traindir = os.path.join(config.DATASET.ROOT + '/images', config.DATASET.TRAIN_SET) valdir = os.path.join(config.DATASET.ROOT + '/images', config.DATASET.TEST_SET) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform_train = transforms.Compose([ transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) transform_valid = transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ]) train_dataset = datasets.ImageFolder(traindir, transform_train) valid_dataset = datasets.ImageFolder(valdir, transform_valid) else: assert dataset_name == "cifar10", "Only CIFAR-10 is supported at this phase" classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # For reference normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) transform_train = transforms.Compose([ transforms.RandomCrop( 32, padding=4), # ** Remove this line if no data augmentation ** transforms.RandomHorizontalFlip( ), # ** Remove this line if no data augmentation ** transforms.ToTensor(), normalize, ]) transform_valid = transforms.Compose([ transforms.ToTensor(), normalize, ]) train_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}', train=True, download=True, transform=transform_train) valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}', train=False, download=True, transform=transform_valid) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) # Learning rate scheduler if lr_scheduler is None: if config.TRAIN.LR_SCHEDULER != 'step': lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, len(train_loader) * config.TRAIN.END_EPOCH, eta_min=1e-6) elif isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) # Training code for epoch in range(last_epoch, config.TRAIN.END_EPOCH): topk = (1, 5) if dataset_name == 'imagenet' else (1, ) if config.TRAIN.LR_SCHEDULER == 'step': lr_scheduler.step() # train for one epoch train(config, train_loader, model, criterion, optimizer, lr_scheduler, epoch, final_output_dir, tb_log_dir, writer_dict, topk=topk) torch.cuda.empty_cache() # evaluate on validation set perf_indicator = validate(config, valid_loader, model, criterion, lr_scheduler, epoch, final_output_dir, tb_log_dir, writer_dict, topk=topk) torch.cuda.empty_cache() writer_dict['writer'].flush() if perf_indicator > best_perf: best_perf = perf_indicator best_model = True else: best_model = False logger.info('=> saving checkpoint to {}'.format(final_output_dir)) save_checkpoint( { 'epoch': epoch + 1, 'model': config.MODEL.NAME, 'state_dict': model.module.state_dict(), 'perf': perf_indicator, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), }, best_model, final_output_dir, filename='checkpoint.pth.tar') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config) dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) logger.info(get_model_summary(model, dump_input)) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) model = torch.nn.DataParallel(model, device_ids=gpus).cuda() # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() optimizer = get_optimizer(config, model) best_perf = 0.0 best_model = False last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_perf = checkpoint['perf'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) best_model = True if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) # Data loading code traindir = os.path.join(config.DATASET.ROOT, config.DATASET.TRAIN_SET) valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) #增加的代码。 #print(train_dataset.classes) #根据分的文件夹的名字来确定的类别 with open("class.txt", "w") as f1: for classname in train_dataset.classes: f1.write(classname + "\n") #print(train_dataset.class_to_idx) #按顺序为这些类别定义索引为0,1... with open("classToIndex.txt", "w") as f2: for key, value in train_dataset.class_to_idx.items(): f2.write(str(key) + " " + str(value) + '\n') #print(train_dataset.imgs) #返回从所有文件夹中得到的图片的路径以及其类别 train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True) valid_loader = torch.utils.data.DataLoader( datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ])), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() # train for one epoch train(config, train_loader, model, criterion, optimizer, epoch, final_output_dir, tb_log_dir, writer_dict) # evaluate on validation set perf_indicator = validate(config, valid_loader, model, criterion, final_output_dir, tb_log_dir, writer_dict) if perf_indicator > best_perf: best_perf = perf_indicator best_model = True else: best_model = False logger.info('=> saving checkpoint to {}'.format(final_output_dir)) save_checkpoint( { 'epoch': epoch + 1, 'model': config.MODEL.NAME, 'state_dict': model.module.state_dict(), 'perf': perf_indicator, 'optimizer': optimizer.state_dict(), }, best_model, final_output_dir, filename='checkpoint.pth.tar') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) final_pth_file = os.path.join(final_output_dir, 'HRNet.pt') print("final_pth_file:", final_pth_file) torch.save(model.module, final_pth_file) writer_dict['writer'].close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'valid') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.' + config.MODEL.NAME + '.get_contrastive_net')(config).cuda() dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])).cuda() logger.info(get_model_summary(model, dump_input)) if config.TEST.MODEL_FILE: logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE)) model.load_state_dict(torch.load(config.TEST.MODEL_FILE)) else: model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info('=> loading model from {}'.format(model_state_file)) model.load_state_dict(torch.load(model_state_file)) gpus = list(config.GPUS) model = torch.nn.DataParallel(model, device_ids=gpus).cuda() # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() # Data loading code valdir = os.path.join(config.DATASET.ROOT + '/images', config.DATASET.TEST_SET) dataset_name = config.DATASET.DATASET if dataset_name == 'imagenet': normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) valid_loader = torch.utils.data.DataLoader( datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ])), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) else: normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) transform_valid = transforms.Compose([ transforms.ToTensor(), normalize, ]) if dataset_name == 'cifar10': valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}', train=False, download=True, transform=transform_valid) elif dataset_name == 'cifar100': valid_dataset = datasets.CIFAR100(root=f'{config.DATASET.ROOT}', train=False, download=True, transform=transform_valid) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) model.eval() embs = [] targets = [] with torch.no_grad(): for i, (input, target) in enumerate(valid_loader): output = model(input, train_step=-1)[-1].cpu().numpy() embs.extend(output) targets.extend(target.cpu().numpy()) embs = np.asarray(embs) print(targets, flush=True) targets = np.asarray(targets) pca = PCA(n_components=64) embs = pca.fit_transform(embs) tsne = TSNE(n_components=2, random_state=0) embs_2d = tsne.fit_transform(embs) target_ids = range(10) plt.figure(figsize=(10, 10)) colors = 'r', 'g', 'b', 'c', 'm', 'yellow', 'k', 'lime', 'orange', 'purple' classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # For reference for i, c, label in zip(target_ids, colors, classes): plt.scatter(embs_2d[targets == i, 0], embs_2d[targets == i, 1], c=c, label=label) plt.legend() plt.savefig(os.path.join(final_output_dir, 'tsne.png'))
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(config) writer_dict = { 'writer': SummaryWriter(tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) distributed = len(gpus) > 1 device = torch.device('cuda:{}'.format(args.local_rank)) # build model model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) if args.local_rank == 0: logger.info(model) tot_params = sum(p.numel() for p in model.parameters()) / 1000000.0 logger.info(f">>> total params: {tot_params:.2f}M") # provide the summary of model dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[0], config.TRAIN.IMAGE_SIZE[1])) logger.info(get_model_summary(model.to(device), dump_input.to(device))) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) if distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group( backend="nccl", init_method="env://", ) # prepare data train_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TRAIN_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=config.TRAIN.MULTI_SCALE, flip=config.TRAIN.FLIP, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TRAIN.BASE_SIZE, crop_size=tuple(config.TRAIN.IMAGE_SIZE), # (height, width) scale_factor=config.TRAIN.SCALE_FACTOR) if distributed: train_sampler = DistributedSampler(train_dataset) else: train_sampler = None trainloader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE and train_sampler is None, num_workers=config.WORKERS, pin_memory=True, drop_last=True, sampler=train_sampler) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=config.TEST.NUM_SAMPLES, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=tuple(config.TEST.IMAGE_SIZE), # (height, width) ) if distributed: test_sampler = DistributedSampler(test_dataset) else: test_sampler = None testloader = torch.utils.data.DataLoader( test_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU, shuffle=False, num_workers=config.WORKERS, pin_memory=True, sampler=test_sampler) # criterion if config.LOSS.USE_OHEM: criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL, weight=train_dataset.class_weights, thresh=config.LOSS.OHEMTHRESH, min_kept=config.LOSS.OHEMKEEP) else: criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL, weight=train_dataset.class_weights) model_state_file = config.MODEL.PRETRAINED logger.info('=> Loading model from {}'.format(model_state_file)) pretrained_dict = torch.load(model_state_file) model_dict = model.state_dict() pretrained_dict = { k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys() } for k, _ in pretrained_dict.items(): logger.info('=> Loading {} from pretrained model'.format(k)) model_dict.update(pretrained_dict) model.load_state_dict(model_dict) model = FullModel(model, criterion) if distributed: model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = model.to(device) if distributed: model = nn.parallel.DistributedDataParallel( model, device_ids=[args.local_rank], output_device=args.local_rank) # optimizer optimizer = get_optimizer(config, model) epoch_iters = np.int(train_dataset.__len__() / config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus)) best_mIoU = 0 last_epoch = 0 if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file, map_location=lambda storage, loc: storage) best_mIoU = checkpoint['best_mIoU'] last_epoch = checkpoint['epoch'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) start = timeit.default_timer() end_epoch = config.TRAIN.END_EPOCH num_iters = config.TRAIN.END_EPOCH * epoch_iters # learning rate scheduler lr_scheduler_dict = { 'optimizer': optimizer, 'milestones': [s * epoch_iters for s in config.TRAIN.LR_STEP], 'gamma': config.TRAIN.LR_FACTOR, 'max_iters': num_iters, 'last_epoch': last_epoch, 'epoch_iters': epoch_iters } lr_scheduler = get_lr_scheduler(config.TRAIN.LR_SCHEDULER, **lr_scheduler_dict) for epoch in range(last_epoch, end_epoch): if distributed: train_sampler.set_epoch(epoch) train(config, epoch, end_epoch, epoch_iters, trainloader, optimizer, lr_scheduler, model, writer_dict, device) valid_loss, mean_IoU = validate(config, testloader, model, writer_dict, device) if args.local_rank == 0: logger.info( '=> saving checkpoint to {}'.format(final_output_dir + '/checkpoint.pth.tar')) torch.save( { 'epoch': epoch + 1, 'best_mIoU': best_mIoU, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict() }, os.path.join(final_output_dir, 'checkpoint.pth.tar')) if mean_IoU > best_mIoU: best_mIoU = mean_IoU torch.save(model.module.state_dict(), os.path.join(final_output_dir, 'best.pth')) msg = f'Loss: {valid_loss:.4f}, MeanIU: {mean_IoU: 4.4f}, \ Best_mIoU: {best_mIoU: 4.4f}' logger.info(msg) if epoch == end_epoch - 1: torch.save(model.module.state_dict(), os.path.join(final_output_dir, 'final_state.pth')) writer_dict['writer'].close() end = timeit.default_timer() logger.info(f'Hours: {np.int((end-start)/3600)}') logger.info('Done!')
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(config) writer_dict = { 'writer': SummaryWriter(tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) distributed = len(gpus) > 1 device = torch.device('cuda:{}'.format(args.local_rank)) # build model model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) if args.local_rank == 0: # provide the summary of model dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.to(device), dump_input.to(device))) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) if distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group( backend="nccl", init_method="env://", ) # prepare data crop_size = (config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]) train_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TRAIN_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=config.TRAIN.MULTI_SCALE, flip=config.TRAIN.FLIP, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TRAIN.BASE_SIZE, crop_size=crop_size, downsample_rate=config.TRAIN.DOWNSAMPLERATE, scale_factor=config.TRAIN.SCALE_FACTOR) if distributed: train_sampler = DistributedSampler(train_dataset) else: train_sampler = None trainloader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE and train_sampler is None, num_workers=config.WORKERS, pin_memory=True, drop_last=True, sampler=train_sampler) if config.DATASET.EXTRA_TRAIN_SET: extra_train_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.EXTRA_TRAIN_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=config.TRAIN.MULTI_SCALE, flip=config.TRAIN.FLIP, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TRAIN.BASE_SIZE, crop_size=crop_size, downsample_rate=config.TRAIN.DOWNSAMPLERATE, scale_factor=config.TRAIN.SCALE_FACTOR) if distributed: extra_train_sampler = DistributedSampler(extra_train_dataset) else: extra_train_sampler = None extra_trainloader = torch.utils.data.DataLoader( extra_train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE and extra_train_sampler is None, num_workers=config.WORKERS, pin_memory=True, drop_last=True, sampler=extra_train_sampler) test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=config.TEST.NUM_SAMPLES, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=test_size, center_crop_test=config.TEST.CENTER_CROP_TEST, downsample_rate=1) if distributed: test_sampler = DistributedSampler(test_dataset) else: test_sampler = None testloader = torch.utils.data.DataLoader( test_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU, shuffle=False, num_workers=config.WORKERS, pin_memory=True, sampler=test_sampler) # criterion if config.LOSS.USE_OHEM: criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL, thres=config.LOSS.OHEMTHRES, min_kept=config.LOSS.OHEMKEEP, weight=train_dataset.class_weights) else: criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL, weight=train_dataset.class_weights) model = FullModel(model, criterion) model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = model.to(device) model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank) # optimizer if config.TRAIN.OPTIMIZER == 'sgd': optimizer = torch.optim.SGD( [{ 'params': filter(lambda p: p.requires_grad, model.parameters()), 'lr': config.TRAIN.LR }], lr=config.TRAIN.LR, momentum=config.TRAIN.MOMENTUM, weight_decay=config.TRAIN.WD, nesterov=config.TRAIN.NESTEROV, ) else: raise ValueError('Only Support SGD optimizer') epoch_iters = np.int(train_dataset.__len__() / config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus)) best_mIoU = 0 last_epoch = 0 if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file, map_location=lambda storage, loc: storage) best_mIoU = checkpoint['best_mIoU'] last_epoch = checkpoint['epoch'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) start = timeit.default_timer() end_epoch = config.TRAIN.END_EPOCH + config.TRAIN.EXTRA_EPOCH num_iters = config.TRAIN.END_EPOCH * epoch_iters extra_iters = config.TRAIN.EXTRA_EPOCH * epoch_iters for epoch in range(last_epoch, end_epoch): if distributed: train_sampler.set_epoch(epoch) if epoch >= config.TRAIN.END_EPOCH: train(config, epoch - config.TRAIN.END_EPOCH, config.TRAIN.EXTRA_EPOCH, epoch_iters, config.TRAIN.EXTRA_LR, extra_iters, extra_trainloader, optimizer, model, writer_dict, device) else: train(config, epoch, config.TRAIN.END_EPOCH, epoch_iters, config.TRAIN.LR, num_iters, trainloader, optimizer, model, writer_dict, device) valid_loss, mean_IoU, IoU_array = validate(config, testloader, model, writer_dict, device) if args.local_rank == 0: logger.info( '=> saving checkpoint to {}'.format(final_output_dir + 'checkpoint.pth.tar')) torch.save( { 'epoch': epoch + 1, 'best_mIoU': best_mIoU, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), }, os.path.join(final_output_dir, 'checkpoint.pth.tar')) if mean_IoU > best_mIoU: best_mIoU = mean_IoU torch.save(model.module.state_dict(), os.path.join(final_output_dir, 'best.pth')) msg = 'Loss: {:.3f}, MeanIU: {: 4.4f}, Best_mIoU: {: 4.4f}'.format( valid_loss, mean_IoU, best_mIoU) logging.info(msg) logging.info(IoU_array) if epoch == end_epoch - 1: torch.save(model.module.state_dict(), os.path.join(final_output_dir, 'final_state.pth')) writer_dict['writer'].close() end = timeit.default_timer() logger.info('Hours: %d' % np.int((end - start) / 3600)) logger.info('Done')
def main(): args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.' + config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: model_state_file = os.path.join( final_output_dir, 'best.pth') # turn final_state.pth to best.pth logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load(model_state_file) if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] model_dict = model.state_dict() pretrained_dict = { k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys() } for k, _ in pretrained_dict.items(): logger.info('=> loading {} from pretrained model'.format(k)) model_dict.update(pretrained_dict) model.load_state_dict(model_dict) gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # prepare data test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=test_size, downsample_rate=1) testloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=config.WORKERS, pin_memory=True) start = timeit.default_timer() if 'val' in config.DATASET.TEST_SET: # mean_IoU, IoU_array, pixel_acc, mean_acc = testval(config, # test_dataset, # testloader, # model) # # msg = 'MeanIU: {: 4.4f}, Pixel_Acc: {: 4.4f}, \ # Mean_Acc: {: 4.4f}, Class IoU: '.format(mean_IoU, # pixel_acc, mean_acc) # logging.info(msg) # logging.info(IoU_array) F, J = myTestval(config, test_dataset, testloader, model) msg = 'F: {: 4.4f}, J: {: 4.4f} '.format(F, J) logging.info(msg) elif 'test' in config.DATASET.TEST_SET: # import pdb # pdb.set_trace() test(config, test_dataset, testloader, model, sv_dir=r'/raid/wj/HRNet-Semantic-Segmentation/results/faceparse') end = timeit.default_timer() logger.info('Mins: %d' % np.int((end - start) / 60)) logger.info('Done')
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(config) writer_dict = { 'writer': SummaryWriter(tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) # build model model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) # prepare data crop_size = (config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]) train_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TRAIN_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=config.TRAIN.MULTI_SCALE, flip=config.TRAIN.FLIP, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TRAIN.BASE_SIZE, crop_size=crop_size, downsample_rate=config.TRAIN.DOWNSAMPLERATE, scale_factor=config.TRAIN.SCALE_FACTOR) trainloader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=True, drop_last=True) test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=config.TEST.NUM_SAMPLES, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=test_size, downsample_rate=1) testloader = torch.utils.data.DataLoader( test_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) # criterion if config.LOSS.USE_OHEM: criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL, thres=config.LOSS.OHEMTHRES, min_kept=config.LOSS.OHEMKEEP, weight=train_dataset.class_weights) else: criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL, weight=train_dataset.class_weights) model = FullModel(model, criterion) model = nn.DataParallel(model, device_ids=gpus).cuda() # optimizer if config.TRAIN.OPTIMIZER == 'sgd': optimizer = torch.optim.SGD( [{ 'params': filter(lambda p: p.requires_grad, model.parameters()), 'lr': config.TRAIN.LR }], lr=config.TRAIN.LR, momentum=config.TRAIN.MOMENTUM, weight_decay=config.TRAIN.WD, nesterov=config.TRAIN.NESTEROV, ) else: raise ValueError('Only Support SGD optimizer') epoch_iters = np.int(train_dataset.__len__() / config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus)) model_state_file = os.path.join(final_output_dir, 'best.pth') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) best_mIoU = checkpoint['best_mIoU'] last_epoch = checkpoint['epoch'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) epoch = 0 valid_loss, mean_IoU = validate(config, testloader, model, writer_dict, epoch) msg = 'Validation Loss: {:.3f}, MeanIU: {: 4.4f}, Best_mIoU: {: 4.4f}'.format( valid_loss, mean_IoU, best_mIoU) logging.info(msg)
def main(): args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'vis_gradcam') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load(model_state_file) model_dict = model.state_dict() pretrained_dict = { k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys() } for k, _ in pretrained_dict.items(): logger.info('=> loading {} from pretrained model'.format(k)) model_dict.update(pretrained_dict) model.load_state_dict(model_dict) # Wrap original model with NBDT if config.NBDT.USE_NBDT: from nbdt.model import SoftSegNBDT model = SoftSegNBDT(config.NBDT.DATASET, model, hierarchy=config.NBDT.HIERARCHY, classes=class_names) device = 'cuda' if torch.cuda.is_available() else 'cpu' model = model.to(device).eval() # Retrieve input image corresponding to args.image_index test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=test_size, downsample_rate=1) # Define target layer as final convolution layer if not specified if args.target_layers: target_layers = args.target_layers.split(',') else: for name, module in list(model.named_modules())[::-1]: if isinstance(module, nn.Conv2d): target_layers = [name] break logger.info('Target layers set to {}'.format(str(target_layers))) # Append model. to target layers if using nbdt if config.NBDT.USE_NBDT: target_layers = ['model.' + layer for layer in target_layers] def generate_and_save_saliency(image_index, pixel_i=None, pixel_j=None, crop_size=None, normalize=False): """too lazy to move out to global lol""" nonlocal maximum, minimum, label # Generate GradCAM + save heatmap heatmaps = [] raw_image = retrieve_raw_image(test_dataset, image_index) should_crop = crop_size is not None and pixel_i is not None and pixel_j is not None if should_crop: raw_image = crop(pixel_i, pixel_j, crop_size, raw_image, is_tensor=False) for layer in target_layers: gradcam_region = gradcam.generate(target_layer=layer, normalize=False) if should_crop: gradcam_region = crop(pixel_i, pixel_j, crop_size, gradcam_region, is_tensor=True) maximum = max(float(gradcam_region.max()), maximum) minimum = min(float(gradcam_region.min()), minimum) logger.info(f'=> Bounds: ({minimum}, {maximum})') heatmaps.append(gradcam_region) output_dir = generate_output_dir(final_output_dir, args.vis_mode, layer, config.NBDT.USE_NBDT, nbdt_node_wnid, args.crop_size, args.nbdt_node_wnids_for) save_path = generate_save_path(output_dir, gradcam_kwargs) logger.info('Saving {} heatmap at {}...'.format( args.vis_mode, save_path)) if normalize: gradcam_region = GradCAM.normalize(gradcam_region) save_gradcam(save_path, gradcam_region, raw_image, save_npy=not args.skip_save_npy) else: save_gradcam(save_path, gradcam_region, raw_image, minimum=minimum, maximum=maximum, save_npy=not args.skip_save_npy) output_dir_original = output_dir + '_original' os.makedirs(output_dir_original, exist_ok=True) save_path_original = generate_save_path(output_dir_original, gradcam_kwargs, ext='jpg') logger.info('Saving {} original at {}...'.format( args.vis_mode, save_path_original)) cv2.imwrite(save_path_original, raw_image) if crop_size and pixel_i and pixel_j: continue output_dir += '_overlap' os.makedirs(output_dir, exist_ok=True) save_path_overlap = generate_save_path(output_dir, gradcam_kwargs, ext='npy') save_path_plot = generate_save_path(output_dir, gradcam_kwargs, ext='jpg') if not args.skip_save_npy: logger.info('Saving {} overlap data at {}...'.format( args.vis_mode, save_path_overlap)) logger.info('Saving {} overlap plot at {}...'.format( args.vis_mode, save_path_plot)) save_overlap(save_path_overlap, save_path_plot, gradcam_region, label, save_npy=not args.skip_save_npy) if len(heatmaps) > 1: combined = torch.prod(torch.stack(heatmaps, dim=0), dim=0) combined /= combined.max() save_path = generate_save_path(final_output_dir, args.vis_mode, gradcam_kwargs, 'combined', config.NBDT.USE_NBDT, nbdt_node_wnid) logger.info('Saving combined {} heatmap at {}...'.format( args.vis_mode, save_path)) save_gradcam(save_path, combined, raw_image) nbdt_node_wnids = args.nbdt_node_wnid or [] cls = args.nbdt_node_wnids_for if cls: assert config.NBDT.USE_NBDT, 'Must be using NBDT' from nbdt.data.custom import Node assert hasattr(model, 'rules') and hasattr(model.rules, 'nodes'), \ 'NBDT must have rules with nodes' logger.info("Getting nodes leading up to class leaf {}...".format(cls)) leaf_to_path_nodes = Node.get_leaf_to_path(model.rules.nodes) cls_index = class_names.index(cls) leaf = model.rules.nodes[0].wnids[cls_index] path_nodes = leaf_to_path_nodes[leaf] nbdt_node_wnids = [ item['node'].wnid for item in path_nodes if item['node'] ] def run(): nonlocal maximum, minimum, label, gradcam_kwargs for image_index in get_image_indices(args.image_index, args.image_index_range): image, label, _, name = test_dataset[image_index] image = torch.from_numpy(image).unsqueeze(0).to(device) logger.info("Using image {}...".format(name)) pred_probs, pred_labels = gradcam.forward(image) maximum, minimum = -1000, 0 logger.info(f'=> Starting bounds: ({minimum}, {maximum})') if args.crop_for and class_names.index(args.crop_for) not in label: print( f'Skipping image {image_index} because no {args.crop_for} found' ) continue if getattr(Saliency, 'whole_image', False): assert not ( args.pixel_i or args.pixel_j or args.pixel_i_range or args.pixel_j_range), \ 'the "Whole" saliency method generates one map for the whole ' \ 'image, not for specific pixels' gradcam_kwargs = {'image': image_index} if args.suffix: gradcam_kwargs['suffix'] = args.suffix gradcam.backward(pred_labels[:, [0], :, :]) generate_and_save_saliency(image_index) if args.crop_size <= 0: continue if args.crop_for: cls_index = class_names.index(args.crop_for) label = torch.Tensor(label).to(pred_labels.device) # is_right_class = pred_labels[0,0,:,:] == cls_index # is_correct = pred_labels == label is_right_class = is_correct = label == cls_index #TODO:tmp pixels = (is_right_class * is_correct).nonzero() pixels = get_random_pixels(args.pixel_max_num_random, pixels, seed=cls_index) else: assert (args.pixel_i or args.pixel_i_range) and (args.pixel_j or args.pixel_j_range) pixels = get_pixels(args.pixel_i, args.pixel_j, args.pixel_i_range, args.pixel_j_range, args.pixel_cartesian_product) logger.info(f'Running on {len(pixels)} pixels.') for pixel_i, pixel_j in pixels: pixel_i, pixel_j = int(pixel_i), int(pixel_j) assert pixel_i < test_size[0] and pixel_j < test_size[1], \ "Pixel ({},{}) is out of bounds for image of size ({},{})".format( pixel_i,pixel_j,test_size[0],test_size[1]) # Run backward pass # Note: Computes backprop wrt most likely predicted class rather than gt class gradcam_kwargs = { 'image': image_index, 'pixel_i': pixel_i, 'pixel_j': pixel_j } if args.suffix: gradcam_kwargs['suffix'] = args.suffix logger.info( f'Running {args.vis_mode} on image {image_index} at pixel ({pixel_i},{pixel_j}). Using filename suffix: {args.suffix}' ) output_pixel_i, output_pixel_j = compute_output_coord( pixel_i, pixel_j, test_size, pred_probs.shape[2:]) if not getattr(Saliency, 'whole_image', False): gradcam.backward(pred_labels[:, [0], :, :], output_pixel_i, output_pixel_j) if args.crop_size <= 0: generate_and_save_saliency(image_index) else: generate_and_save_saliency(image_index, pixel_i, pixel_j, args.crop_size) logger.info(f'=> Final bounds are: ({minimum}, {maximum})') # Instantiate wrapper once, outside of loop Saliency = METHODS[args.vis_mode] gradcam = Saliency(model=model, candidate_layers=target_layers, use_nbdt=config.NBDT.USE_NBDT, nbdt_node_wnid=None) maximum, minimum, label, gradcam_kwargs = -1000, 0, None, {} for nbdt_node_wnid in nbdt_node_wnids: if config.NBDT.USE_NBDT: logger.info("Using logits from node with wnid {}...".format( nbdt_node_wnid)) gradcam.set_nbdt_node_wnid(nbdt_node_wnid) run() if not nbdt_node_wnids: nbdt_node_wnid = None run()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED dataset = config.DATASET.DATASET if dataset.startswith('imagenet'): model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config) else: model = eval('models.' + config.MODEL.NAME + '.get_cls_net_cifar')(config) dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])) logger.info(get_model_summary(model, dump_input)) # n_flops, n_params = measure_model(model, 32, 32, torch.zeros(1,3,32,32)) # logger.info("param size = %fMB", n_params[0]/1e6) # logger.info("flops = %fM", n_flops[0]/1e6) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) model = torch.nn.DataParallel(model, device_ids=gpus).cuda() # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() optimizer = get_optimizer(config, model) best_perf = 0.0 best_model = False last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_perf = checkpoint['perf'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) best_model = True if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) # Data loading code if dataset.startswith('imagenet'): traindir = os.path.join(config.DATASET.ROOT, config.DATASET.TRAIN_SET) valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET) normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True) valid_loader = torch.utils.data.DataLoader( datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ])), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) else: CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124] CIFAR_STD = [0.24703233, 0.24348505, 0.26158768] if config.TRAIN.AUGMENT == 'autoaugment': print("==>use autoaugment") train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4, fill=128), transforms.RandomHorizontalFlip(), CIFAR10Policy(), transforms.ToTensor(), Cutout_v2(n_holes=1, length=16), transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ]) else: train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ]) if config.TRAIN.AUGMENT == 'cutout': train_transform.transforms.append(Cutout(16)) valid_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ]) train_dataset = datasets.CIFAR10(root=config.DATASET.ROOT, train=True, download=True, transform=train_transform) valid_dataset = datasets.CIFAR10(root=config.DATASET.ROOT, train=False, download=True, transform=valid_transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() # train for one epoch train(config, train_loader, model, criterion, optimizer, epoch, final_output_dir, tb_log_dir, writer_dict) # evaluate on validation set perf_indicator = validate(config, valid_loader, model, criterion, final_output_dir, tb_log_dir, writer_dict) if perf_indicator > best_perf: best_perf = perf_indicator best_model = True else: best_model = False logger.info('=> saving checkpoint to {}'.format(final_output_dir)) save_checkpoint( { 'epoch': epoch + 1, 'model': config.MODEL.NAME, 'state_dict': model.module.state_dict(), 'perf': perf_indicator, 'optimizer': optimizer.state_dict(), }, best_model, final_output_dir, filename='checkpoint.pth.tar') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225] args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.' + config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: # model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth') model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load( '/home/kong/Documents/DDRNet.Pytorch/DDRNet.Pytorch/output/face/ddrnet23_slim/checkpoint.pth.tar' ) if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] # print(pretrained_dict.keys()) new_state = { k: v for k, v in pretrained_dict.items() if k in model.state_dict() } model.load_state_dict(new_state) model = model.cuda() torch.save(model.state_dict(), 'model_best_bacc.pth.tar', _use_new_zipfile_serialization=False) # gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() # prepare data test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) print(test_size) img = cv2.imread( "/home/kong/Documents/DDRNet.Pytorch/DDRNet.Pytorch/images/418cd0c0b416d93bc5a129834537f2e1.jpeg" ) img = cv2.resize(img, (512, 512)) image = img.astype(np.float32)[:, :, ::-1] image = image / 255.0 image -= mean image /= std image = image.transpose((2, 0, 1)) image = torch.from_numpy(image) # image = image.permute((2, 0, 1)) # print(image.shape) image = image.unsqueeze(0) image = image.cuda() start = time.time() # print("input : ",image) for i in range(1): out = model(image) end = time.time() #print("Cuda 1000 images inference time : ",1000.0/(end - start)) outadd = out[0] * 1.0 + out[1] * 0.4 out0 = out[0].squeeze(dim=0) out1 = out[1].squeeze(dim=0) # print(out0.size(),out0[0,1,1],out0[1,1,1]) # print("out:",out0) outadd = outadd.squeeze(dim=0) out0 = F.softmax(out0, dim=0) out1 = F.softmax(out1, dim=0) outadd = F.softmax(outadd, dim=0) out0 = torch.argmax(out0, dim=0) out1 = torch.argmax(out1, dim=0) outadd = torch.argmax(outadd, dim=0) pred0 = out0.detach().cpu().numpy() pred1 = out1.detach().cpu().numpy() predadd = outadd.detach().cpu().numpy() pred0 = pred0 * 255 pred1 = pred1 * 255 predadd = predadd * 255 ####================= alpha channel =========================# print("pred0:", pred0.shape, img.shape) pred0 = np.array(pred0, np.uint8) pred0up = cv2.resize(pred0, (512, 512)) png = np.dstack((img, pred0up)) bg = cv2.imread( "/home/kong/Downloads/4aeb26a89778f73261ccef283e70992f.jpeg") addpng = merge_img(bg, png, 500, 1012, 100, 612) cv2.imwrite("png.png", addpng) pred_ch = np.zeros(pred0.shape) pred_rgb0 = np.array([pred_ch, pred_ch, pred0]) pred_rgb1 = np.array([pred_ch, pred_ch, pred1]) pred_rgbadd = np.array([predadd, pred_ch, predadd]) pred_rgb0 = pred_rgb0.transpose(1, 2, 0) pred_rgb1 = pred_rgb1.transpose(1, 2, 0) pred_rgbadd = pred_rgbadd.transpose(1, 2, 0) pred_rgb0 = cv2.resize(pred_rgb0, (img.shape[1], img.shape[0])) pred_rgb1 = cv2.resize(pred_rgb1, (img.shape[1], img.shape[0])) pred_rgbadd = cv2.resize(pred_rgbadd, (img.shape[1], img.shape[0])) dst = cv2.addWeighted(img, 0.7, pred_rgb0.astype(np.uint8), 0.3, 0) dst1 = cv2.addWeighted(img, 0.7, pred_rgb1.astype(np.uint8), 0.3, 0) dstadd = cv2.addWeighted(img, 0.7, pred_rgbadd.astype(np.uint8), 0.3, 0) imgadd = np.hstack((img, dstadd)) cv2.imwrite("a242.jpg", imgadd)
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') model_func = 'models.{0}.get_{0}'.format(config.MODEL.NAME) model = eval(model_func)(config, pretrained=True) dump_input = torch.rand((1, 3, config.MODEL.IMAGE_SIZE[0], config.MODEL.IMAGE_SIZE[1])) print(get_model_summary(model, dump_input)) gpus = list(config.GPUS) device = torch.device('cuda:{}'.format(gpus[0]) if torch.cuda.is_available() else 'cpu') model.to(device) optimizer = get_optimizer(config, model) if config.FP16: model, optimizer = amp.initialize(model, optimizer, opt_level='O1') criterion = {'reid': nn.CrossEntropyLoss(), 'sr': nn.MSELoss()} if config.TRAIN.SR_FILTER: criterion['sr'] = MSEFilterLoss() best_perf = 0.0 best_model = False last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_perf = checkpoint['perf'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) if config.FP16: amp.load_state_dict(checkpoint['amp']) logger.info("=> loaded checkpoint (epoch {})" .format(checkpoint['epoch'])) best_model = True writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } train_dataset = ImagePersonDataset(config, mode='train') train_sampler = RandomIdentitySampler( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), num_instances=4 ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), sampler=train_sampler, num_workers=config.WORKERS, pin_memory=True ) valid_loader = torch.utils.data.DataLoader( ImagePersonDataset(config, mode='val'), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), num_workers=config.WORKERS, pin_memory=True ) since = time.time() for epoch in range(last_epoch, config.TRAIN.END_EPOCH): print('Epoch {}/{}'.format(epoch + 1, config.TRAIN.END_EPOCH)) train(config, train_loader, model, criterion, optimizer, epoch, device, writer_dict) perf_indicator = validate(config, valid_loader, model, criterion, device, writer_dict) if perf_indicator > best_perf: best_perf = perf_indicator best_model = True else: best_model = False logger.info('=> saving checkpoint to {}'.format(final_output_dir)) checkpoint = { 'epoch': epoch + 1, 'model': config.MODEL.NAME, 'state_dict': model.state_dict(), 'perf': perf_indicator, 'optimizer': optimizer.state_dict() } if config.FP16: checkpoint['amp'] = amp.state_dict() save_checkpoint(checkpoint, best_model, final_output_dir, filename='checkpoint.pth.tar') h, m, s = get_hms_from_sec(time.time() - since) print('=> total training time: {}h {}m {}s'.format(h, m, s)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info('saving final model state to {}'.format( final_model_state_file)) torch.save(model.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225] args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.' + config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: # model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth') model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load( '/home/hwits/Documents/CarVid/DDRNet/DDRNet.pytorch/model_best_bacc.pth.tar' ) if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] newstate_dict = { k: v for k, v in pretrained_dict.items() if k in model.state_dict() } # print(pretrained_dict.keys()) model.load_state_dict(newstate_dict) model = model.cuda() print(model) x = torch.ones((1, 3, 1024, 1024)).cuda() model_trt = torch2trt(model, [x]) y = model(x) print(len(y)) y_trt = model_trt(x) print(y_trt.shape) print(len(y_trt)) # check the output against PyTorch print(torch.max(torch.abs(y - y_trt))) with open('ddrnet.engine', "wb") as f: f.write(model_trt.engine.serialize())
def main(): args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'vis_erf') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load(model_state_file) model_dict = model.state_dict() pretrained_dict = { k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys() } for k, _ in pretrained_dict.items(): logger.info('=> loading {} from pretrained model'.format(k)) model_dict.update(pretrained_dict) model.load_state_dict(model_dict) device = 'cuda' if torch.cuda.is_available() else 'cpu' model = model.to(device).eval() # Generate input logger.info('Generating input of ones...') input = torch.ones((1, 3, config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])).float().to(device) input = torch.autograd.Variable(input, requires_grad=True) # Run input through network + compute backward pass logger.info('Computing backward pass...') output = model(input) output_grad = torch.zeros_like(output) center_i, center_j = output.shape[2] // 2, output.shape[3] // 2 erf_center_i, erf_center_j = center_i + args.offset_i, center_j + args.offset_j output_grad[0, 0, erf_center_i, erf_center_j] = 1. output.backward(gradient=output_grad) input_grad = input.grad[0, 0].cpu().data.numpy() # Compute receptive field rectangle nonzero_i = np.nonzero(np.sum(input_grad, axis=1)) min_i, max_i = np.min(nonzero_i), np.max(nonzero_i) nonzero_j = np.nonzero(np.sum(input_grad, axis=0)) min_j, max_j = np.min(nonzero_j), np.max(nonzero_j) rect_origin = (min_j, min_i) rect_h, rect_w = max_i - min_i, max_j - min_j rect = patches.Rectangle(rect_origin, rect_w, rect_h, linewidth=1, edgecolor='g', facecolor='none') # Save plot of gradient logger.info('Saving gradient map...') fig = plt.figure() ax = fig.add_subplot(111) ax.imshow(input_grad, cmap='coolwarm') ax.add_patch(rect) title_values = [ output.shape[2], output.shape[3], erf_center_i, erf_center_j, rect_origin[1], rect_origin[0], rect_h, rect_w ] ax.set_title( "Output Size: {} x {}, Output Pixel: ({},{})\n ERF Origin: ({},{}), ERF Size: {} x {}" .format(*title_values)) save_path = os.path.join( final_output_dir, 'erf_{}_{}.png'.format(args.offset_i, args.offset_j)) plt.savefig(save_path)
def main(): mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] args = parse_args() logger, final_output_dir, _ = create_logger( config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.'+config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.'+config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]) ) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: # model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth') model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load('/home/hwits/Documents/CarVid/DDRNet/segmentation/best_val_smaller.pth') # print("pretrained_dict:",pretrained_dict.keys()) if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] # print(pretrained_dict.keys()) new_state = {k[6:]:v for k,v in pretrained_dict.items() if k[6:] in model.state_dict()} model.load_state_dict(new_state) model = model.cuda() torch.save(model.state_dict(), 'model_best_bacc.pth.tar', _use_new_zipfile_serialization=False) # gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() # prepare data test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) print(test_size) img = cv2.imread("/home/hwits/berlin_000528_000019_leftImg8bit.png") img = cv2.resize(img,(1024,1024)) image = img.astype(np.float32)[:, :, ::-1] image = image / 255.0 image -= mean image /= std image = image.transpose((2,0,1)) image = torch.from_numpy(image) # image = image.permute((2, 0, 1)) print(image.shape) image = image.unsqueeze(0) image = image.cuda() start = time.time() # print("input : ",image) for i in range(1): out= model(image) end = time.time() print("out : ",out[0].shape) pred = F.interpolate( out[0], (img.shape[0],img.shape[1]), mode='nearest' ) print("results : ",pred.shape) _, pred = torch.max(pred, dim=1) pred = pred.squeeze(0).cpu().numpy() map16.visualize_result(img, pred, '.', 'cityscape.jpg') print("Cuda 1000 images inference time : ",1000.0/(end - start)) outadd = out[0]*1.0 + out[1]*0.4 out0 = out[0].squeeze(dim=0) out1 = out[1].squeeze(dim=0) outadd = outadd.squeeze(dim=0) out0 = F.softmax(out0,dim=0) out1 = F.softmax(out1,dim=0) outadd = F.softmax(outadd,dim=0) out0 = torch.argmax(out0,dim=0) out1 = torch.argmax(out1,dim=0) outadd = torch.argmax(outadd,dim=0) pred0 = out0.detach().cpu().numpy() pred1 = out1.detach().cpu().numpy() predadd = outadd.detach().cpu().numpy() pred0 = pred0*255 pred1 = pred1*255 predadd = predadd*255 pred_ch = np.zeros(pred0.shape) pred_rgb0 = np.array([pred_ch,pred_ch,pred0]) pred_rgb1 = np.array([pred_ch,pred_ch,pred1]) pred_rgbadd = np.array([predadd,pred_ch,predadd]) pred_rgb0 = pred_rgb0.transpose(1,2,0) pred_rgb1 = pred_rgb1.transpose(1,2,0) pred_rgbadd = pred_rgbadd.transpose(1,2,0) pred_rgb0 = cv2.resize(pred_rgb0,(img.shape[1],img.shape[0])) pred_rgb1 = cv2.resize(pred_rgb1,(img.shape[1],img.shape[0])) pred_rgbadd = cv2.resize(pred_rgbadd,(img.shape[1],img.shape[0])) dst=cv2.addWeighted(img,0.7,pred_rgb0.astype(np.uint8),0.3,0) dst1=cv2.addWeighted(img,0.7,pred_rgb1.astype(np.uint8),0.3,0) dstadd=cv2.addWeighted(img,0.7,pred_rgbadd.astype(np.uint8),0.3,0) imgadd = np.vstack((img,pred_rgb0,dst,pred_rgb1, dst1,pred_rgbadd, dstadd)) cv2.imwrite("a22.jpg",imgadd)
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'valid') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED if config.MODEL.DOWNSAMPLE: model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config).cuda() else: model = eval('models.' + config.MODEL.NAME + '.get_linear_net')(config).cuda() dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])).cuda() logger.info(get_model_summary(model, dump_input)) if config.TEST.MODEL_FILE: logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE)) model.load_state_dict(torch.load(config.TEST.MODEL_FILE)) else: model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info('=> loading model from {}'.format(model_state_file)) model.load_state_dict(torch.load(model_state_file)) gpus = list(config.GPUS) model = torch.nn.DataParallel(model, device_ids=gpus).cuda() # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() # Data loading code valdir = os.path.join(config.DATASET.ROOT + '/images', config.DATASET.TEST_SET) dataset_name = config.DATASET.DATASET if dataset_name == 'imagenet': normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) valid_loader = torch.utils.data.DataLoader( datasets.ImageFolder( valdir, transforms.Compose([ transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)), transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]), transforms.ToTensor(), normalize, ])), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) else: normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) transform_valid = transforms.Compose([ transforms.ToTensor(), normalize, ]) if dataset_name == 'cifar10': valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}', train=False, download=True, transform=transform_valid) elif dataset_name == 'cifar100': valid_dataset = datasets.CIFAR100(root=f'{config.DATASET.ROOT}', train=False, download=True, transform=transform_valid) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) # evaluate on validation set validate(config, valid_loader, model, criterion, None, -1, final_output_dir, tb_log_dir, None, topk=(1, 5))
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = create_logger( config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = config.CUDNN.ENABLED model = eval('models.'+config.MODEL.NAME+'.get_cls_net')( config) dump_input = torch.rand( (1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]) ) logger.info(get_model_summary(model, dump_input)) # copy model file this_dir = os.path.dirname(__file__) models_dst_dir = os.path.join(final_output_dir, 'models') if os.path.exists(models_dst_dir): shutil.rmtree(models_dst_dir) shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir) writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = [0,1] model = torch.nn.DataParallel(model, device_ids=gpus).cuda() # define loss function (criterion) and optimizer criterion = torch.nn.CrossEntropyLoss().cuda() optimizer = get_optimizer(config, model) best_perf = 0.0 best_model = False last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'checkpoint.pth.tar') if os.path.isfile(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_perf = checkpoint['perf'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger.info("=> loaded checkpoint (epoch {})" .format(checkpoint['epoch'])) best_model = True if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch-1 ) else: lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch-1 ) # Data loading code traindir = os.path.join(config.DATASET.ROOT, config.DATASET.TRAIN_SET) valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET) train_dataset = datasets(dataset_root='./Data/',split='train',size= config.MODEL.IMAGE_SIZE[0]) test_dataset = datasets(dataset_root='./Data/',split='test',size= config.MODEL.IMAGE_SIZE[0]) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True ) valid_loader = torch.utils.data.DataLoader( test_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus), shuffle=True, num_workers=config.WORKERS, pin_memory=True ) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() # train for one epoch train(config, train_loader, model, criterion, optimizer, epoch, final_output_dir, tb_log_dir, writer_dict) # evaluate on validation set perf_indicator = validate(config, valid_loader, model, criterion, final_output_dir, tb_log_dir, writer_dict) if perf_indicator > best_perf: best_perf = perf_indicator best_model = True torch.save(best_model, './best-model.pth') else: best_model = False logger.info('=> saving checkpoint to {}'.format(final_output_dir)) save_checkpoint({ 'epoch': epoch + 1, 'model': config.MODEL.NAME, 'state_dict': model.module.state_dict(), 'perf': perf_indicator, 'optimizer': optimizer.state_dict(), }, best_model, final_output_dir, filename='checkpoint.pth.tar') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth.tar') logger.info('saving final model state to {}'.format( final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): args = parse_args() logger, final_output_dir, _ = create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model if torch.__version__.startswith('1'): module = eval('models.' + config.MODEL.NAME) module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: # model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth') model_state_file = os.path.join(final_output_dir, 'best.pth') logger.info('=> loading model from {}'.format(model_state_file)) pretrained_dict = torch.load(model_state_file) if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict'] model_dict = model.state_dict() pretrained_dict = { k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys() } for k, _ in pretrained_dict.items(): logger.info('=> loading {} from pretrained model'.format(k)) model_dict.update(pretrained_dict) model.load_state_dict(model_dict) model = model.cuda() # gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() # prepare data test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=test_size, downsample_rate=1) testloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=config.WORKERS, pin_memory=True) start = timeit.default_timer() test(config, test_dataset, testloader, model, sv_dir=final_output_dir + '/test_result') end = timeit.default_timer() logger.info('Mins: %d' % np.int((end - start) / 60)) logger.info('Done')
def main(): args = parse_args() logger, final_output_dir, _ = create_logger( config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # build model model = models.gsto_hrnet.get_seg_model(config) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]) ) logger.info(get_model_summary(model.cuda(), dump_input.cuda())) if config.TEST.MODEL_FILE: model_state_file = config.TEST.MODEL_FILE else: model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info('=> loading model from {}'.format(model_state_file)) model_state_file = config.TEST.MODEL_FILE pretrained_dict = torch.load(model_state_file) model_dict = model.state_dict() pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()} model_dict.update(pretrained_dict) model.load_state_dict(model_dict) gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # prepare data test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]) test_dataset = eval('datasets.' + config.DATASET.DATASET)( root=config.DATASET.ROOT, list_path=config.DATASET.TEST_SET, num_samples=None, num_classes=config.DATASET.NUM_CLASSES, multi_scale=False, flip=False, ignore_label=config.TRAIN.IGNORE_LABEL, base_size=config.TEST.BASE_SIZE, crop_size=test_size, downsample_rate=1) testloader = torch.utils.data.DataLoader( test_dataset, batch_size=1, shuffle=False, num_workers=config.WORKERS, pin_memory=True) start = timeit.default_timer() if 'val' in config.DATASET.TEST_SET: mean_IoU, IoU_array, pixel_acc, mean_acc = testval(config, test_dataset, testloader, model) msg = 'MeanIU: {: 4.4f}, Pixel_Acc: {: 4.4f}, \ Mean_Acc: {: 4.4f}, Class IoU: '.format(mean_IoU, pixel_acc, mean_acc) logging.info(msg) logging.info(IoU_array) elif 'test' in config.DATASET.TEST_SET: test(config, test_dataset, testloader, model, sv_dir=final_output_dir) end = timeit.default_timer() logger.info('Mins: %d' % np.int((end - start) / 60)) logger.info('Done')
def __init__(self, config): super().__init__(config) logger.info(config) # cudnn related setting cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.deterministic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) self.writer_dict = { 'writer': SummaryWriter(config.SUMMARY_DIR), 'train_global_steps': 0, 'valid_global_steps': 0, } self.checkpoint_dir = config.CHECKPOINT_DIR self.output_dir = config.OUTPUT_DIR self.num_classes = config.DATASET.NUM_CLASSES if config.MODEL.NAME == 'unet': self.model = UNet(3, config.DATASET.NUM_CLASSES) elif config.MODEL.NAME == 'dl_resnet': self.model = DeepLabv3_plus_res( 3, n_classes=config.DATASET.NUM_CLASSES, os=config.MODEL.OS, pretrained=True) elif config.MODEL.NAME == 'dl_xception': self.model = DeepLabv3_plus_xce( 3, n_classes=config.DATASET.NUM_CLASSES, os=config.MODEL.OS, pretrained=True) dump_input = torch.rand( (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])) logger.info(get_model_summary(self.model.cuda(), dump_input.cuda())) self.train_dataset = RssraiDataset('train', config.DATASET, mean=[0.2797, 0.3397, 0.3250], std=[0.1271, 0.1446, 0.1320]) self.trainloader = torch.utils.data.DataLoader( self.train_dataset, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=True, drop_last=True) self.test_dataset = RssraiDataset('val', config.DATASET, mean=[0.2797, 0.3397, 0.3250], std=[0.1271, 0.1446, 0.1320]) self.testloader = torch.utils.data.DataLoader( self.test_dataset, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=True) if config.LOSS.USE_OHEM: self.criterion = OhemCrossEntropy(thres=config.LOSS.OHEMTHRES, min_kept=config.LOSS.OHEMKEEP) else: self.criterion = CrossEntropy() if len(gpus) > 1: self.model = nn.DataParallel(self.model, device_ids=gpus).cuda() else: self.model.cuda() # optimizer if config.TRAIN.OPTIMIZER == 'sgd': self.optimizer = torch.optim.SGD( [{ 'params': filter(lambda p: p.requires_grad, self.model.parameters()), 'lr': config.TRAIN.LR }], lr=config.TRAIN.LR, momentum=config.TRAIN.MOMENTUM, weight_decay=config.TRAIN.WD, nesterov=config.TRAIN.NESTEROV, ) elif config.TRAIN.OPTIMIZER == 'adam': self.optimizer = torch.optim.Adam([{ 'params': filter(lambda p: p.requires_grad, self.model.parameters()), 'lr': config.TRAIN.LR }], lr=config.TRAIN.LR, weight_decay=config.TRAIN.WD) else: raise ValueError('Only Support SGD optimizer') self.lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( self.optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR) self.start_epoch = 0 self.current_epoch = self.start_epoch self.end_epoch = config.TRAIN.EPOCH init_model_weights(self.model) if config.TRAIN.AUTO_RESUME: self.load_checkpoint( os.path.join(config.CHECKPOINT_DIR, 'checkpoint.pth.tar')) self.best_mIoU = -1