def setup_pose_net(config, prepared, **kwargs): """ Create a pose network Parameters ---------- config : CfgNode Network configuration prepared : bool True if the network has been prepared before kwargs : dict Extra parameters for the network Returns ------- pose_net : nn.Module Created pose network """ print0(pcolor('PoseNet: %s' % config.name, 'yellow')) pose_net = load_class_args_create( config.name, paths=[ 'fbnet.networks.pose', ], args={ **config, **kwargs }, ) if not prepared and config.checkpoint_path is not '': pose_net = load_network(pose_net, config.checkpoint_path, ['pose_net', 'pose_network']) return pose_net
def setup_depth_net2(config, prepared, **kwargs): """ Create a depth network Parameters ---------- config : CfgNode Network configuration prepared : bool True if the network has been prepared before kwargs : dict Extra parameters for the network Returns ------- depth_net : nn.Module Create depth network """ print0(pcolor('DepthNet2: %s' % config.name, 'yellow')) depth_net = load_class_args_create( config.name, paths=[ 'fbnet.networks.depth2', ], args={ **config, **kwargs }, ) if not prepared and config.checkpoint_path is not '': depth_net = load_network(depth_net, config.checkpoint_path, ['depth_net2', 'disp_network2']) return depth_net
def prepare_model(self, resume=None): """Prepare self.model (incl. loading previous state)""" print0(pcolor('### Preparing Model', 'green')) self.model = setup_model(self.config.model, self.config.prepared) # Resume model if available if resume: print0( pcolor('### Resuming from {}'.format(resume['file']), 'magenta', attrs=['bold'])) self.model = load_network(self.model, resume['state_dict'], 'model') if 'epoch' in resume: self.current_epoch = resume['epoch']
def setup_model(config, prepared, **kwargs): """ Create a model Parameters ---------- config : CfgNode Model configuration (cf. configs/default_config.py) prepared : bool True if the model has been prepared before kwargs : dict Extra parameters for the model Returns ------- model : nn.Module Created model """ print0(pcolor('Model: %s' % config.name, 'yellow')) model = load_class(config.name, paths=[ 'fbnet.models', ])(**{ **config.loss, **kwargs }) # Add depth network if required if model.network_requirements['depth_net']: model.add_depth_net(setup_depth_net(config.depth_net, prepared)) # Add depth network if required if model.network_requirements['depth_net2']: model.add_depth_net2(setup_depth_net2(config.depth_net2, prepared)) # Add pose network if required if model.network_requirements['pose_net']: model.add_pose_net(setup_pose_net(config.pose_net, prepared)) # If a checkpoint is provided, load pretrained model if not prepared and config.checkpoint_path is not '': model = load_network(model, config.checkpoint_path, 'model') # Return model return model