config.merge_from_list(args.opts)
config.freeze()
save_dir = os.path.join(config.save_dir)
mkdir(save_dir)
logger = setup_logger("inference", save_dir, 0)
logger.info("Running with config:\n{}".format(config))
device = torch.device(config.device)
num_types = len(config.boundaries) + 2
generator = Generator(BertConfig(type_vocab_size=num_types))
generator = generator.to(device)
g_checkpointer = Checkpointer(model=generator, logger=logger)
g_checkpointer.load(config.model_path, True)
dataset = COCOCaptionDataset(root=config.data_dir,
split='test',
boundaries=config.boundaries)
data_loader = make_data_loader(dataset=dataset,
collate_fn=collate_fn_infer,
batch_size=config.samples_per_gpu,
num_workers=config.num_workers,
split='test')
pred_dict = inference(generator, data_loader, device)
logger.info(f"Saving results to {save_dir}/caption_results.json")
with open(os.path.join(save_dir, 'caption_results.json'), 'w') as f:
json.dump(pred_dict, f)
generator.load_weights(config.pretrained_bert)
else:
extra_checkpoint_data = checkpointer.load(config.model_path)
arguments.update(extra_checkpoint_data)
dataset = COCOCaptionDataset(
root=config.data_dir,
split='trainrestval',
boundaries=config.boundaries,
)
data_loader = make_data_loader(
dataset=dataset,
collate_fn=collate_fn_train,
batch_size=config.samples_per_gpu,
num_workers=config.num_workers,
max_iter=config.scheduler.max_steps,
split='trainrestval',
is_distributed=config.distributed,
start_iter=arguments['iteration'],
)
if config.distributed:
generator = DistributedDataParallel(
module=generator,
device_ids=[args.local_rank],
output_device=args.local_rank,
)
train(generator=generator,
optimizer=optimizer,
data_loader=data_loader,