def build_caption_model(caption_config: Dict, cuda_device: torch.device):
"""
Parameters
----------
caption_config : Dict
Dict of BUTD and Detectron model configuration.
cuda_device : torch.device
Torch device to load the model to.
Returns
-------
(model, caption_processor, text_processor) : List[object]
Returns the model, caption and text processor
"""
with open(caption_config["butd_model"]["config_yaml"]) as f:
butd_config = yaml.load(f, Loader=yaml.FullLoader)
butd_config = ConfigNode(butd_config)
butd_config.training_parameters.evalai_inference = True
registry.register("config", butd_config)
caption_processor, text_processor = init_processors(
caption_config, butd_config)
if cuda_device == torch.device('cpu'):
state_dict = torch.load(caption_config["butd_model"]["model_pth"],
map_location='cpu')
else:
state_dict = torch.load(caption_config["butd_model"]["model_pth"])
model_config = butd_config.model_attributes.butd
model_config.model_data_dir = caption_config["model_data_dir"]
model = BUTD(model_config)
model.build()
model.init_losses_and_metrics()
if list(state_dict.keys())[0].startswith('module') and \
not hasattr(model, 'module'):
state_dict = multi_gpu_state_to_single(state_dict)
model.load_state_dict(state_dict)
model.to(cuda_device)
model.eval()
return model, caption_processor, text_processor
def _build_pythia_model(self):
state_dict = torch.load('model_data/butd.pth')
model_config = self.config.model_attributes.butd
model_config.model_data_dir = "../../"
model = BUTD(model_config)
model.build()
model.init_losses_and_metrics()
if list(state_dict.keys())[0].startswith('module') and \
not hasattr(model, 'module'):
state_dict = self._multi_gpu_state_to_single(state_dict)
model.load_state_dict(state_dict)
model.to("cuda")
model.eval()
return model
def _build_pythia_model(self):
state_dict = torch.load(model_pth)
model_config = self.config.model_attributes.butd
model_config.model_data_dir = "content/"
model = BUTD(model_config)
model.build()
model.init_losses_and_metrics()
if list(state_dict.keys())[0].startswith('module') and \
not hasattr(model, 'module'):
state_dict = self._multi_gpu_state_to_single(state_dict)
model.load_state_dict(state_dict)
# print("DEBUG Model's state_dict:")
# for param_tensor in model.state_dict():
# print("DEBUG:", param_tensor, "\t", model.state_dict()[param_tensor].size())
model.to("cuda")
model.eval()
return model
def _build_model(self):
state_dict = torch.load(butd_checkpoint)
model_config = self.config.model_attributes.butd
# specify the root directory of pre-trained model
model_config.model_data_dir = os.path.dirname(
os.path.abspath('__file__'))
if self.use_constrained:
model_config.inference.type = 'constrained_beam_search'
model = BUTD(model_config)
model.build()
model.init_losses_and_metrics()
if list(state_dict.keys())[0].startswith('module') and not hasattr(
model, 'module'):
state_dict = self._multi_gpu_state_to_single(state_dict)
# load pre-trained state dict
model.load_state_dict(state_dict)
model.to(device)
model.eval()
return model