def demo(parser): tools.add_arguments_for_module(parser, models, argument_for_class='model', default='FlowNet2') args = parser.parse_args() if args.number_gpus < 0: args.number_gpus = torch.cuda.device_count() kwargs = tools.kwargs_from_args(args, 'model') model = tools.module_to_dict(models)[args.model] model = model(args, **kwargs) args.cuda = not args.no_cuda and torch.cuda.is_available() print('Initializing CUDA') model = model.cuda() print("Loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) print("Loaded checkpoint '{}' (at epoch {})".format( args.resume, checkpoint['epoch'])) print("Initializing save directory: {}".format(args.save)) if not os.path.exists(args.save): os.makedirs(args.save) stats = inference(args=args, model=model) print("Demo test done")
def get_config(self): json_file = "config.json" with open(json_file, 'r') as config_file: config_dict = json.load(config_file) # convert the dictionary to a namespace using bunch lib self.args = Bunch(config_dict) self.args.name_model = self.name_model self.args.model = self.name_model.replace('-', '') self.args.resume = os.path.join(self.args.dir_checkpoints,(self.args.name_model + "_checkpoint.pth.tar")) self.args.model_class = tools.module_to_dict(source.models)[self.args.model] self.args.loss_class = tools.module_to_dict(source.losses)[self.args.loss] self.args.cuda = not self.args.no_cuda and torch.cuda.is_available() self.args.inference_dir = os.path.join("{}/inference".format(self.args.save), self.args.name_model)
# Parse the official arguments with tools.TimerBlock("Parsing Arguments") as block: args = parser.parse_args() # Get argument defaults (hastag #thisisahack) parser.add_argument('--IGNORE', action='store_true') defaults = vars(parser.parse_args(['--IGNORE'])) # Print all arguments, color the non-defaults for argument, value in sorted(vars(args).items()): reset = colorama.Style.RESET_ALL color = reset if value == defaults[ argument] else colorama.Fore.MAGENTA block.log('{}{}: {}{}'.format(color, argument, value, reset)) args.model_class = tools.module_to_dict(models)[args.model] args.optimizer_class = tools.module_to_dict( torch.optim)[args.optimizer] args.loss_class = tools.module_to_dict(losses)[args.loss] args.training_dataset_class = tools.module_to_dict(datasets)[ args.training_dataset] args.validation_dataset_class = tools.module_to_dict(datasets)[ args.validation_dataset] args.inference_dataset_class = tools.module_to_dict(datasets)[ args.inference_dataset] args.cuda = not args.no_cuda and torch.cuda.is_available() args.current_hash = subprocess.check_output( ["git", "rev-parse", "HEAD"]).rstrip() args.log_file = join(args.save, 'args.txt')
# Parse the official arguments with tools.TimerBlock("Parsing Arguments") as block: args = parser.parse_args() if args.number_gpus < 0 : args.number_gpus = torch.cuda.device_count() # Get argument defaults (hastag #thisisahack) parser.add_argument('--IGNORE', action='store_true') defaults = vars(parser.parse_args(['--IGNORE'])) # Print all arguments, color the non-defaults for argument, value in sorted(vars(args).items()): reset = colorama.Style.RESET_ALL color = reset if value == defaults[argument] else colorama.Fore.MAGENTA block.log('{}{}: {}{}'.format(color, argument, value, reset)) args.model_class = tools.module_to_dict(models)[args.model] args.inference_dataset_class = tools.module_to_dict(datasets)[args.inference_dataset] args.cuda = not args.no_cuda and torch.cuda.is_available() # args.current_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).rstrip() args.log_file = join(args.save, 'args.txt') # dict to collect activation gradients (for training debug purpose) args.grads = {} args.total_epochs = 1 args.inference_dir = "{}/inference".format(args.save) print('Source Code') # print((' Current Git Hash: {}\n'.format(args.current_hash)))
parser.add_argument( '--fp16', action='store_true', help='Run model in pseudo-fp16 mode (fp16 storage fp32 math).') parser.add_argument("--rgb_max", type=float, default=255.) tools.add_arguments_for_module(parser, models, argument_for_class='model', default='FlowNet2S') # Parse the official arguments with tools.TimerBlock("Parsing Arguments") as block: args = parser.parse_args() args.model_class = tools.module_to_dict(models)[args.model] # Dynamically load model with parameters passed in via "--model_[param]=[value]" arguments with tools.TimerBlock("Building {} model".format(args.model)) as block: kwargs = tools.kwargs_from_args(args, 'model') model = args.model_class(args, **kwargs) block.log('Number of parameters: {}'.format( sum([ p.data.nelement() if p.requires_grad else 0 for p in model.parameters() ]))) # Load weights if needed, otherwise randomly initialize if args.resume and os.path.isfile(args.resume):
def initialize_args(): if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--start_epoch', type=int, default=1) parser.add_argument('--total_epochs', type=int, default=10000) parser.add_argument('--batch_size', '-b', type=int, default=8, help="Batch size") parser.add_argument( '--train_n_batches', type=int, default=-1, help= 'Number of min-batches per epoch. If < 0, it will be determined by training_dataloader' ) parser.add_argument( '--crop_size', type=int, nargs='+', default=[256, 256], help="Spatial dimension to crop training samples for training") parser.add_argument('--gradient_clip', type=float, default=None) parser.add_argument('--schedule_lr_frequency', type=int, default=0, help='in number of iterations (0 for no schedule)') parser.add_argument('--schedule_lr_fraction', type=float, default=10) parser.add_argument("--rgb_max", type=float, default=255.) parser.add_argument('--number_workers', '-nw', '--num_workers', type=int, default=8) parser.add_argument('--number_gpus', '-ng', type=int, default=-1, help='number of GPUs to use') parser.add_argument('--no_cuda', action='store_true') parser.add_argument('--seed', type=int, default=1) parser.add_argument('--name', default='run', type=str, help='a name to append to the save directory') parser.add_argument('--save', '-s', default='./work', type=str, help='directory for saving') parser.add_argument('--validation_frequency', type=int, default=5, help='validate every n epochs') parser.add_argument('--validation_n_batches', type=int, default=-1) parser.add_argument( '--render_validation', action='store_true', help= 'run inference (save flows to file) and every validation_frequency epoch' ) parser.add_argument('--inference', action='store_true') parser.add_argument( '--inference_size', type=int, nargs='+', default=[-1, -1], help= 'spatial size divisible by 64. default (-1,-1) - largest possible valid size would be used' ) parser.add_argument('--inference_batch_size', type=int, default=1) parser.add_argument('--inference_n_batches', type=int, default=-1) parser.add_argument('--save_flow', action='store_true', help='save predicted flows to file') parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--log_frequency', '--summ_iter', type=int, default=1, help="Log every n batches") parser.add_argument('--skip_training', action='store_true') parser.add_argument('--skip_validation', action='store_true') parser.add_argument( '--fp16', action='store_true', help='Run model in pseudo-fp16 mode (fp16 storage fp32 math).') parser.add_argument( '--fp16_scale', type=float, default=1024., help= 'Loss scaling, positive power of 2 values can improve fp16 convergence.' ) tools.add_arguments_for_module(parser, models, argument_for_class='model', default='FlowNet2') tools.add_arguments_for_module(parser, losses, argument_for_class='loss', default='L1Loss') tools.add_arguments_for_module(parser, torch.optim, argument_for_class='optimizer', default='Adam', skip_params=['params']) tools.add_arguments_for_module( parser, datasets, argument_for_class='training_dataset', default='MpiSintelFinal', skip_params=['is_cropped'], parameter_defaults={'root': './MPI-Sintel/flow/training'}) tools.add_arguments_for_module(parser, datasets, argument_for_class='validation_dataset', default='MpiSintelClean', skip_params=['is_cropped'], parameter_defaults={ 'root': './MPI-Sintel/flow/training', 'replicates': 1 }) tools.add_arguments_for_module(parser, datasets, argument_for_class='inference_dataset', default='MpiSintelClean', skip_params=['is_cropped'], parameter_defaults={ 'root': './MPI-Sintel/flow/training', 'replicates': 1 }) main_dir = os.path.dirname(os.path.realpath(__file__)) os.chdir(main_dir) # Parse the official arguments with tools.TimerBlock("Parsing Arguments") as block: args = parser.parse_args() if args.number_gpus < 0: args.number_gpus = torch.cuda.device_count() # Get argument defaults (hastag #thisisahack) parser.add_argument('--IGNORE', action='store_true') defaults = vars(parser.parse_args(['--IGNORE'])) # Print all arguments, color the non-defaults for argument, value in sorted(vars(args).items()): reset = colorama.Style.RESET_ALL color = reset if value == defaults[ argument] else colorama.Fore.MAGENTA block.log('{}{}: {}{}'.format(color, argument, value, reset)) args.model_class = tools.module_to_dict(models)[args.model] args.optimizer_class = tools.module_to_dict( torch.optim)[args.optimizer] args.loss_class = tools.module_to_dict(losses)[args.loss] args.training_dataset_class = tools.module_to_dict(datasets)[ args.training_dataset] args.validation_dataset_class = tools.module_to_dict(datasets)[ args.validation_dataset] args.inference_dataset_class = tools.module_to_dict(datasets)[ args.inference_dataset] args.cuda = not args.no_cuda and torch.cuda.is_available() args.current_hash = subprocess.check_output( ["git", "rev-parse", "HEAD"]).rstrip() args.log_file = join(args.save, 'args.txt') # dict to collect activation gradients (for training debug purpose) args.grads = {} if args.inference: args.skip_validation = True args.skip_training = True args.total_epochs = 1 args.inference_dir = "{}/inference".format(args.save) print('Source Code') print((' Current Git Hash: {}\n'.format(args.current_hash))) # Change the title for `top` and `pkill` commands setproctitle.setproctitle(args.save) # Dynamically load the dataset class with parameters passed in via "--argument_[param]=[value]" arguments with tools.TimerBlock("Initializing Datasets") as block: args.effective_batch_size = args.batch_size * args.number_gpus args.effective_inference_batch_size = args.inference_batch_size * args.number_gpus args.effective_number_workers = args.number_workers * args.number_gpus gpuargs = { 'num_workers': args.effective_number_workers, 'pin_memory': True, 'drop_last': True } if args.cuda else {} inf_gpuargs = gpuargs.copy() inf_gpuargs['num_workers'] = args.number_workers if exists(args.training_dataset_root): train_dataset = args.training_dataset_class( args, True, **tools.kwargs_from_args(args, 'training_dataset')) block.log('Training Dataset: {}'.format(args.training_dataset)) block.log('Training Input: {}'.format(' '.join([ str([d for d in x.size()]) for x in train_dataset[0][0] ]))) block.log('Training Targets: {}'.format(' '.join([ str([d for d in x.size()]) for x in train_dataset[0][1] ]))) train_loader = DataLoader(train_dataset, batch_size=args.effective_batch_size, shuffle=True, **gpuargs) if exists(args.validation_dataset_root): validation_dataset = args.validation_dataset_class( args, True, **tools.kwargs_from_args(args, 'validation_dataset')) block.log('Validation Dataset: {}'.format( args.validation_dataset)) block.log('Validation Input: {}'.format(' '.join([ str([d for d in x.size()]) for x in validation_dataset[0][0] ]))) block.log('Validation Targets: {}'.format(' '.join([ str([d for d in x.size()]) for x in validation_dataset[0][1] ]))) validation_loader = DataLoader( validation_dataset, batch_size=args.effective_batch_size, shuffle=False, **gpuargs) if exists(args.inference_dataset_root): inference_dataset = args.inference_dataset_class( args, False, **tools.kwargs_from_args(args, 'inference_dataset')) block.log('Inference Dataset: {}'.format( args.inference_dataset)) block.log('Inference Input: {}'.format(' '.join([ str([d for d in x.size()]) for x in inference_dataset[0][0] ]))) block.log('Inference Targets: {}'.format(' '.join([ str([d for d in x.size()]) for x in inference_dataset[0][1] ]))) inference_loader = DataLoader( inference_dataset, batch_size=args.effective_inference_batch_size, shuffle=False, **inf_gpuargs) # Dynamically load model and loss class with parameters passed in via "--model_[param]=[value]" or "--loss_[param]=[value]" arguments with tools.TimerBlock("Building {} model".format(args.model)) as block: class ModelAndLoss(nn.Module): def __init__(self, args): super(ModelAndLoss, self).__init__() kwargs = tools.kwargs_from_args(args, 'model') self.model = args.model_class(args, **kwargs) kwargs = tools.kwargs_from_args(args, 'loss') self.loss = args.loss_class(args, **kwargs) def forward(self, data, target, inference=False): output = self.model(data) loss_values = self.loss(output, target) if not inference: return loss_values else: return loss_values, output model_and_loss = ModelAndLoss(args) block.log('Effective Batch Size: {}'.format( args.effective_batch_size)) block.log('Number of parameters: {}'.format( sum([ p.data.nelement() if p.requires_grad else 0 for p in model_and_loss.parameters() ]))) # assing to cuda or wrap with dataparallel, model and loss if args.cuda and (args.number_gpus > 0) and args.fp16: block.log('Parallelizing') model_and_loss = nn.parallel.DataParallel( model_and_loss, device_ids=list(range(args.number_gpus))) block.log('Initializing CUDA') model_and_loss = model_and_loss.cuda().half() torch.cuda.manual_seed(args.seed) param_copy = [ param.clone().type(torch.cuda.FloatTensor).detach() for param in model_and_loss.parameters() ] elif args.cuda and args.number_gpus > 0: block.log('Initializing CUDA') model_and_loss = model_and_loss.cuda() block.log('Parallelizing') model_and_loss = nn.parallel.DataParallel( model_and_loss, device_ids=list(range(args.number_gpus))) torch.cuda.manual_seed(args.seed) else: block.log('CUDA not being used') torch.manual_seed(args.seed) # Load weights if needed, otherwise randomly initialize if args.resume and os.path.isfile(args.resume): block.log("Loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) if not args.inference: args.start_epoch = checkpoint['epoch'] best_err = checkpoint['best_EPE'] model_and_loss.module.model.load_state_dict( checkpoint['state_dict']) block.log("Loaded checkpoint '{}' (at epoch {})".format( args.resume, checkpoint['epoch'])) elif args.resume and args.inference: block.log("No checkpoint found at '{}'".format(args.resume)) quit() else: block.log("Random initialization") block.log("Initializing save directory: {}".format(args.save)) if not os.path.exists(args.save): os.makedirs(args.save) train_logger = SummaryWriter(log_dir=os.path.join( args.save, 'train'), comment='training') validation_logger = SummaryWriter(log_dir=os.path.join( args.save, 'validation'), comment='validation') # Dynamically load the optimizer with parameters passed in via "--optimizer_[param]=[value]" arguments with tools.TimerBlock("Initializing {} Optimizer".format( args.optimizer)) as block: kwargs = tools.kwargs_from_args(args, 'optimizer') if args.fp16: optimizer = args.optimizer_class( [p for p in param_copy if p.requires_grad], **kwargs) else: optimizer = args.optimizer_class([ p for p in model_and_loss.parameters() if p.requires_grad ], **kwargs) for param, default in list(kwargs.items()): block.log("{} = {} ({})".format(param, default, type(default))) # Log all arguments to file for argument, value in sorted(vars(args).items()): block.log2file(args.log_file, '{}: {}'.format(argument, value)) return args
def infer_flownet(in_path, out_path, reverse, downscale_factor=1): args = SimpleNamespace(model="FlowNet2", reverse=reverse, downscale_factor=downscale_factor, start_epoch=1, total_epochs=10000, batch_size=8, train_n_batches=-1, crop_size=[256, 256], gradient_clip=None, schedule_lr_frequency=0, schedule_lr_fraction=10, rgb_max=255., number_workers=8, number_gpus=-1, no_cuda=False, seed=1, name='run', in_path=in_path, save=out_path, validation_frequency=5, validation_n_batches=-1, render_validation=False, inference=True, inference_visualize=True, inference_size=[-1, -1], inference_batch_size=1, inference_n_batches=-1, save_flow=True, resume='./FlowNet2_checkpoint.pth.tar', log_frequency=1, skip_training=False, skip_validation=False, fp16=False, fp16_scale=1024., loss='L1Loss', optimizer='Adam', training_dataset='MpiSintelFinal', root='./MPI-Sintel/flow/training', validation_dataset='MpiSintelClean', inference_dataset='MpiSintelClean', IGNORE=False) print(args.in_path) main_dir = os.path.dirname(os.path.realpath(__file__)) os.chdir(main_dir) # Parse the official arguments with tools.TimerBlock("Parsing Arguments") as block: if args.number_gpus < 0: args.number_gpus = torch.cuda.device_count() args.model_class = tools.module_to_dict(models)[args.model] args.optimizer_class = tools.module_to_dict( torch.optim)[args.optimizer] args.loss_class = tools.module_to_dict(losses)[args.loss] args.cuda = not args.no_cuda and torch.cuda.is_available() args.current_hash = subprocess.check_output( ["git", "rev-parse", "HEAD"]).rstrip() # dict to collect activation gradients (for training debug purpose) args.grads = {} if args.inference: args.skip_validation = True args.skip_training = True args.total_epochs = 1 args.inference_dir = "{}/inference".format(args.save) print('Source Code') print((' Current Git Hash: {}\n'.format(args.current_hash))) # Change the title for `top` and `pkill` commands setproctitle.setproctitle(args.save) # Dynamically load the dataset class with parameters passed in via "--argument_[param]=[value]" arguments with tools.TimerBlock("Initializing Datasets") as block: args.effective_batch_size = args.batch_size * args.number_gpus args.effective_inference_batch_size = args.inference_batch_size * args.number_gpus args.effective_number_workers = args.number_workers * args.number_gpus gpuargs = { 'num_workers': args.effective_number_workers, 'pin_memory': True, 'drop_last': True } if args.cuda else {} inf_gpuargs = gpuargs.copy() inf_gpuargs['num_workers'] = args.number_workers inference_dataset = datasets.ImagesFromFolder(args, is_cropped=False, is_reversed=args.reverse, root=args.in_path) block.log('Inference Input: {}'.format(' '.join( [str([d for d in x.size()]) for x in inference_dataset[0][0]]))) block.log('Inference Targets: {}'.format(' '.join( [str([d for d in x.size()]) for x in inference_dataset[0][1]]))) inference_loader = DataLoader( inference_dataset, batch_size=args.effective_inference_batch_size, shuffle=False, **inf_gpuargs) # Dynamically load model and loss class with parameters passed in via "--model_[param]=[value]" or "--loss_[param]=[value]" arguments with tools.TimerBlock("Building {} model".format(args.model)) as block: class ModelAndLoss(nn.Module): def __init__(self, args): super(ModelAndLoss, self).__init__() kwargs = tools.kwargs_from_args(args, 'model') self.model = args.model_class(args, **kwargs) kwargs = tools.kwargs_from_args(args, 'loss') self.loss = args.loss_class(args, **kwargs) def forward(self, data, target, inference=False): output = self.model(data) loss_values = self.loss(output, target) if not inference: return loss_values else: return loss_values, output model_and_loss = ModelAndLoss(args) # assing to cuda or wrap with dataparallel, model and loss if args.cuda and (args.number_gpus > 0) and args.fp16: model_and_loss = nn.parallel.DataParallel( model_and_loss, device_ids=list(range(args.number_gpus))) model_and_loss = model_and_loss.cuda().half() torch.cuda.manual_seed(args.seed) elif args.cuda and args.number_gpus > 0: model_and_loss = model_and_loss.cuda() model_and_loss = nn.parallel.DataParallel( model_and_loss, device_ids=list(range(args.number_gpus))) torch.cuda.manual_seed(args.seed) else: block.log('CUDA not being used') torch.manual_seed(args.seed) # Load weights if needed, otherwise randomly initialize if args.resume and os.path.isfile(args.resume): checkpoint = torch.load(args.resume) if not args.inference: args.start_epoch = checkpoint['epoch'] model_and_loss.module.model.load_state_dict( checkpoint['state_dict']) elif args.resume and args.inference: block.log("No checkpoint found at '{}'".format(args.resume)) quit() else: block.log("Random initialization") if not os.path.exists(args.save): os.makedirs(args.save) # Reusable function for inference def inference(args, data_loader, model, offset=0): model.eval() if args.save_flow or args.render_validation: flow_folder = out_path # "./output/flo_rev" if args.reverse else "./output/flo" if not os.path.exists(flow_folder): os.makedirs(flow_folder) # visualization folder if args.inference_visualize: flow_vis_folder = out_path + "/" + "png/" if not os.path.exists(flow_vis_folder): os.makedirs(flow_vis_folder) args.inference_n_batches = np.inf if args.inference_n_batches < 0 else args.inference_n_batches progress = tqdm(data_loader, ncols=100, total=np.minimum(len(data_loader), args.inference_n_batches), desc='Inferencing ', leave=True, position=offset) statistics = [] total_loss = 0 ph, pw = inference_dataset.ph, inference_dataset.pw for batch_idx, (data, target) in enumerate(progress): if args.cuda: data, target = [d.cuda(non_blocking=True) for d in data ], [t.cuda(non_blocking=True) for t in target] data, target = [Variable(d) for d in data], [Variable(t) for t in target] # when ground-truth flows are not available for inference_dataset, # the targets are set to all zeros. thus, losses are actually L1 or L2 norms of compute optical flows, # depending on the type of loss norm passed in with torch.no_grad(): losses, output = model(data[0], target[0], inference=True) losses = [torch.mean(loss_value) for loss_value in losses] loss_val = losses[0] # Collect first loss for weight update total_loss += loss_val.item() loss_values = [v.item() for v in losses] statistics.append(loss_values) # import IPython; IPython.embed() if args.save_flow or args.render_validation: for i in range(args.inference_batch_size): _pflow = output[i].data.cpu().numpy().transpose(1, 2, 0) if ph != 0: _pflow = _pflow[ph:-ph, :, :] if pw != 0: _pflow = _pflow[:, pw:-pw, :] flow_utils.writeFlow( join( flow_folder, '%06d.flo' % (batch_idx * args.inference_batch_size + i)), _pflow) # You can comment out the plt block in visulize_flow_file() for real-time visualization if args.inference_visualize: flow_utils.visulize_flow_file( join( flow_folder, '%06d.flo' % (batch_idx * args.inference_batch_size + i)), flow_vis_folder) progress.update(1) if batch_idx == (args.inference_n_batches - 1): break progress.close() return progress = tqdm(list(range(args.start_epoch, args.total_epochs + 1)), miniters=1, ncols=100, desc='Overall Progress', leave=True, position=0) offset = 1 for _ in progress: inference(args=args, data_loader=inference_loader, model=model_and_loss, offset=offset) offset += 1 print("\n")