def run_dag(self,
inputs: TensorDict,
verbose: bool = False,
**kwargs) -> TensorDict:
"""Connects and runs submodules of dag.
Args:
inputs: A dictionary of input tensors fed to the dag.
verbose: Print out dag routing when running.
**kwargs: Other kwargs to pass to submodules, such as keras kwargs.
Returns:
A nested dictionary of all the output tensors.
"""
# Initialize the outputs with inputs to the dag.
outputs = {'inputs': inputs}
# TODO(jesseengel): Remove this cluttering of the base namespace. Only there
# for backwards compatability.
outputs.update(inputs)
# Run through the DAG nodes in sequential order.
for node in self.dag:
# The first element of the node can be either a module or module_key.
module_key, input_keys = node[0], node[1]
module = getattr(self, module_key)
# Optionally specify output keys if module does not return dict.
output_keys = node[2] if len(node) > 2 else None
# Get the inputs to the node.
inputs = [core.nested_lookup(key, outputs) for key in input_keys]
# Duck typing to avoid dealing with multiple inheritance of Group modules.
if is_processor(module):
# Processor modules.
module_outputs = module(*inputs,
return_outputs_dict=True,
**kwargs)
elif is_loss(module):
# Loss modules.
module_outputs = module.get_losses_dict(*inputs, **kwargs)
else:
# Network modules.
module_outputs = module(*inputs, **kwargs)
if not isinstance(module_outputs, dict):
module_outputs = core.to_dict(module_outputs, output_keys)
# Add module outputs to the dictionary.
outputs[module_key] = module_outputs
# Alias final module output as dag output.
# 'out' is a reserved key for final dag output.
outputs['out'] = module_outputs
return outputs
def run_dag(self,
inputs: TensorDict,
verbose: bool = True,
**kwargs) -> TensorDict:
"""Connects and runs submodules of dag.
Args:
inputs: A dictionary of input tensors fed to the dag.
verbose: Print out dag routing when running.
**kwargs: Other kwargs to pass to submodules, such as keras kwargs.
Returns:
A nested dictionary of all the output tensors.
"""
# Initialize the outputs with inputs to the dag.
outputs = {'inputs': inputs}
# TODO(jesseengel): Remove this cluttering of the base namespace. Only there
# for backwards compatability.
outputs.update(inputs)
# Run through the DAG nodes in sequential order.
for node in self.dag:
# The first element of the node can be either a module or module_key.
module_key, input_keys = node[0], node[1]
module = getattr(self, module_key)
# Optionally specify output keys if module does not return dict.
output_keys = node[2] if len(node) > 2 else None
# Get the inputs to the node.
inputs = [core.nested_lookup(key, outputs) for key in input_keys]
if verbose:
shape = lambda d: tf.nest.map_structure(
lambda x: list(x.shape), d)
logging.info('Input to Module: %s\nKeys: %s\nIn: %s\n',
module_key, input_keys, shape(inputs))
if is_processor(module):
# Processor modules.
module_outputs = module(*inputs,
return_outputs_dict=True,
**kwargs)
else:
# Network modules.
module_outputs = module(*inputs, **kwargs)
if not isinstance(module_outputs, dict):
module_outputs = core.to_dict(module_outputs, output_keys)
if verbose:
logging.info('Output from Module: %s\nOut: %s\n', module_key,
shape(module_outputs))
# Add module outputs to the dictionary.
outputs[module_key] = module_outputs
# Alias final module output as dag output.
# 'out' is a reserved key for final dag output.
outputs['out'] = module_outputs
return outputs