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
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