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