def __call__(self, *inputs, **kwargs): """Wrap the layer's __call__() with dictionary inputs and outputs. IMPORTANT: If no input_keys are provided to the constructor, they are inferred from the argument names in call(). If no output_keys are provided to the constructor, they are inferred from return annotation of call() (a list of strings). Example: ``` def call(self, f0_hz, loudness) -> ['amps', 'frequencies']: ... return amps, frequencies ``` Will infer `self.input_keys = ['f0_hz', 'loudness']` and `self.output_keys = ['amps', 'frequencies']`. If input_keys, or output_keys are provided to the constructor they will override these inferred values. Example Usage: The the example above works with both tensor inputs `layer(f0_hz, loudness)` or a dictionary of tensors `layer({'f0_hz':..., 'loudness':...})`, and in both cases will return a dictionary of tensors `{'amps':..., 'frequencies':...}`. Args: *inputs: Arguments passed on to call(). If any arguments are dicts, they will be merged and self.input_keys will be read out of them and passed to call() while other args will be ignored. **kwargs: Keyword arguments passed on to call(). Returns: outputs: A dictionary of layer outputs from call(). If the layer call() returns a dictionary it will be returned directly, otherwise the output tensors will be wrapped in a dictionary {output_key: output_tensor}. """ # Merge all dictionaries provided in inputs. input_dict = {} for v in inputs: if isinstance(v, dict): input_dict.update(v) # If any dicts provided, lookup input tensors from those dicts. # Otherwise, just use inputs list as input tensors. if input_dict: inputs = [core.nested_lookup(key, input_dict) for key in self.input_keys] # Run input tensors through the model. outputs = super().__call__(*inputs, **kwargs) # Return dict if call() returns it. if isinstance(outputs, dict): return outputs # Otherwise make a dict from output_keys. else: outputs = core.make_iterable(outputs) if len(self.output_keys) != len(outputs): raise ValueError(f'Output keys ({self.output_keys}) must have the same' f'length as outputs ({outputs})') return dict(zip(self.output_keys, outputs))
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 get_controls(self, dag_inputs: TensorDict) -> TensorDict: """Run the DAG and get complete outputs dictionary for the processor_group. Args: dag_inputs: A dictionary of input tensors fed to the signal processing processor_group. Returns: A nested dictionary of all the output tensors. """ # Initialize the outputs with inputs to the processor_group. outputs = dag_inputs # Run through the DAG nodes in sequential order. for node in self.dag: # Get the node processor and keys to the node input. processor, keys = node # Logging, only on the first call. if not self.built: logging.info('Connecting node (%s):', processor.name) for i, key in enumerate(keys): logging.info('Input %d: %s', i, key) # Get the inputs to the node. inputs = [core.nested_lookup(key, outputs) for key in keys] # Build the processor only if called the first time in a @tf.function. # Need to explicitly build because we use get_controls() and get_signal() # seperately, (to get intermediates) rather than directly using call(). if not processor.built: processor.build([tensor.shape for tensor in inputs]) # Run processor. controls = processor.get_controls(*inputs) signal = processor.get_signal(**controls) # Add outputs to the dictionary. outputs[processor.name] = {'controls': controls, 'signal': signal} # Get output signal from last processor. output_name = self.processors[-1].name outputs[self.name] = {'signal': outputs[output_name]['signal']} # Logging, only on the first call. if not self.built: logging.info('ProcessorGroup output node (%s)', output_name) return outputs
def get_outputs(self, dag_inputs: TensorDict) -> TensorDict: """Run the DAG and get complete outputs dictionary for the processor_group. Args: dag_inputs: A dictionary of input tensors fed to the signal processing processor_group. Returns: outputs: A nested dictionary of all the output tensors. """ # Initialize the outputs with inputs to the processor_group. outputs = dag_inputs # Run through the DAG nodes in sequential order. for node in self.dag: # Get the node processor and keys to the node input. processor, keys = node # Logging. logging.info('Connecting node (%s):', processor.name) for i, key in enumerate(keys): logging.info('Input %d: %s', i, key) # Get the inputs to the node. inputs = [core.nested_lookup(key, outputs) for key in keys] # Build the processor (does nothing if not the first time). processor.build([tensor.shape for tensor in inputs]) # Run processor. controls = processor.get_controls(*inputs) signal = processor.get_signal(**controls) # Add outputs to the dictionary. outputs[processor.name] = {'controls': controls, 'signal': signal} # Get output signal from last processor. output_name = self.processors[-1].name outputs[self.name] = {'signal': outputs[output_name]['signal']} logging.info('ProcessorGroup output node (%s)', output_name) return outputs
def __call__(self, *inputs, **kwargs): """Wrap the layer's __call__() with dictionary inputs and outputs. IMPORTANT: If no input_keys are provided to the constructor, they are inferred from the argument names in call(). If no output_keys are provided to the constructor, they are inferred from return annotation of call() (a list of strings). Example: ======== ``` def call(self, f0_hz, loudness, power=None) -> ['amps', 'frequencies']: ... return amps, frequencies ``` Will infer `self.input_keys = ['f0_hz', 'loudness']` and `self.output_keys = ['amps', 'frequencies']`. If input_keys, or output_keys are provided to the constructor they will override these inferred values. It will also infer `self.default_input_keys = ['power']`, which it will try to look up the inputs, but use the default values and not throw an error if the key is not in the input dictionary. Example Usage: ============== The the example above works with both tensor inputs `layer(f0_hz, loudness)` or `layer(f0_hz, loudness, power)` or a dictionary of tensors `layer({'f0_hz':..., 'loudness':...})`, or `layer({'f0_hz':..., 'loudness':..., 'power':...})` and in both cases will return a dictionary of tensors `{'amps':..., 'frequencies':...}`. Args: *inputs: Arguments passed on to call(). If any arguments are dicts, they will be merged and self.input_keys will be read out of them and passed to call() while other args will be ignored. **kwargs: Keyword arguments passed on to call(). Returns: outputs: A dictionary of layer outputs from call(). If the layer call() returns a dictionary it will be returned directly, otherwise the output tensors will be wrapped in a dictionary {output_key: output_tensor}. """ # Construct a list of input tensors equal in length and order to the `call` # input signature. # -- Start first with any tensor arguments. # -- Then lookup tensors from input dictionaries. # -- Use default values if not found. # Start by merging all dictionaries of tensors from the input. input_dict = {} for v in inputs: if isinstance(v, dict): input_dict.update(v) # And then strip all dictionaries from the input. inputs = [v for v in inputs if not isinstance(v, dict)] # Add any tensors from kwargs. for key in self.all_input_keys: if key in kwargs: input_dict[key] = kwargs[key] # And strip from kwargs. kwargs = { k: v for k, v in kwargs.items() if k not in self.all_input_keys } # Look up further inputs from the dictionaries. for key in self.input_keys: try: # If key is present use the input_dict value. inputs.append(core.nested_lookup(key, input_dict)) except KeyError: # Skip if not present. pass # Add default arguments. for key, value in zip(self.default_input_keys, self.default_input_values): try: # If key is present, use the input_dict value. inputs.append(core.nested_lookup(key, input_dict)) except KeyError: # Otherwise use the default value if not supplied as non-dict input. if len(inputs) < self.n_inputs: inputs.append(value) # Run input tensors through the model. if len(inputs) != self.n_inputs: raise TypeError( f'{len(inputs)} input tensors extracted from inputs' '(including default args) but the layer expects ' f'{self.n_inputs} tensors.\n' f'Input keys: {self.input_keys}\n' f'Default keys: {self.default_input_keys}\n' f'Default values: {self.default_input_values}\n' f'Input dictionaries: {input_dict}\n' f'Input Tensors (Args, Dicts, and Defaults): {inputs}\n') outputs = super().__call__(*inputs, **kwargs) # Return dict if call() returns it. if isinstance(outputs, dict): return outputs # Otherwise make a dict from output_keys. else: outputs = core.make_iterable(outputs) if len(self.output_keys) != len(outputs): raise ValueError( f'Output keys ({self.output_keys}) must have the same' f'length as outputs ({outputs})') return dict(zip(self.output_keys, outputs))
def _check_tensor_outputs(self, strings_to_check, outputs): for tensor_string in strings_to_check: tensor = core.nested_lookup(tensor_string, outputs) self.assertIsInstance(tensor, (np.ndarray, tf.Tensor))
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