| ## @package rnn_cell |
| # Module caffe2.python.rnn_cell |
| |
| |
| |
| |
| |
| import functools |
| import inspect |
| import logging |
| import numpy as np |
| import random |
| from future.utils import viewkeys |
| |
| from caffe2.proto import caffe2_pb2 |
| from caffe2.python.attention import ( |
| apply_dot_attention, |
| apply_recurrent_attention, |
| apply_regular_attention, |
| apply_soft_coverage_attention, |
| AttentionType, |
| ) |
| from caffe2.python import core, recurrent, workspace, brew, scope, utils |
| from caffe2.python.modeling.parameter_sharing import ParameterSharing |
| from caffe2.python.modeling.parameter_info import ParameterTags |
| from caffe2.python.modeling.initializers import Initializer |
| from caffe2.python.model_helper import ModelHelper |
| |
| |
| def _RectifyName(blob_reference_or_name): |
| if blob_reference_or_name is None: |
| return None |
| if isinstance(blob_reference_or_name, str): |
| return core.ScopedBlobReference(blob_reference_or_name) |
| if not isinstance(blob_reference_or_name, core.BlobReference): |
| raise Exception("Unknown blob reference type") |
| return blob_reference_or_name |
| |
| |
| def _RectifyNames(blob_references_or_names): |
| if blob_references_or_names is None: |
| return None |
| return [_RectifyName(i) for i in blob_references_or_names] |
| |
| |
| class RNNCell(object): |
| ''' |
| Base class for writing recurrent / stateful operations. |
| |
| One needs to implement 2 methods: apply_override |
| and get_state_names_override. |
| |
| As a result base class will provice apply_over_sequence method, which |
| allows you to apply recurrent operations over a sequence of any length. |
| |
| As optional you could add input and output preparation steps by overriding |
| corresponding methods. |
| ''' |
| def __init__(self, name=None, forward_only=False, initializer=None): |
| self.name = name |
| self.recompute_blobs = [] |
| self.forward_only = forward_only |
| self._initializer = initializer |
| |
| @property |
| def initializer(self): |
| return self._initializer |
| |
| @initializer.setter |
| def initializer(self, value): |
| self._initializer = value |
| |
| def scope(self, name): |
| return self.name + '/' + name if self.name is not None else name |
| |
| def apply_over_sequence( |
| self, |
| model, |
| inputs, |
| seq_lengths=None, |
| initial_states=None, |
| outputs_with_grads=None, |
| ): |
| if initial_states is None: |
| with scope.NameScope(self.name): |
| if self.initializer is None: |
| raise Exception("Either initial states " |
| "or initializer have to be set") |
| initial_states = self.initializer.create_states(model) |
| |
| preprocessed_inputs = self.prepare_input(model, inputs) |
| step_model = ModelHelper(name=self.name, param_model=model) |
| input_t, timestep = step_model.net.AddScopedExternalInputs( |
| 'input_t', |
| 'timestep', |
| ) |
| utils.raiseIfNotEqual( |
| len(initial_states), len(self.get_state_names()), |
| "Number of initial state values provided doesn't match the number " |
| "of states" |
| ) |
| states_prev = step_model.net.AddScopedExternalInputs(*[ |
| s + '_prev' for s in self.get_state_names() |
| ]) |
| states = self._apply( |
| model=step_model, |
| input_t=input_t, |
| seq_lengths=seq_lengths, |
| states=states_prev, |
| timestep=timestep, |
| ) |
| |
| external_outputs = set(step_model.net.Proto().external_output) |
| for state in states: |
| if state not in external_outputs: |
| step_model.net.AddExternalOutput(state) |
| |
| if outputs_with_grads is None: |
| outputs_with_grads = [self.get_output_state_index() * 2] |
| |
| # states_for_all_steps consists of combination of |
| # states gather for all steps and final states. It looks like this: |
| # (state_1_all, state_1_final, state_2_all, state_2_final, ...) |
| states_for_all_steps = recurrent.recurrent_net( |
| net=model.net, |
| cell_net=step_model.net, |
| inputs=[(input_t, preprocessed_inputs)], |
| initial_cell_inputs=list(zip(states_prev, initial_states)), |
| links=dict(zip(states_prev, states)), |
| timestep=timestep, |
| scope=self.name, |
| forward_only=self.forward_only, |
| outputs_with_grads=outputs_with_grads, |
| recompute_blobs_on_backward=self.recompute_blobs, |
| ) |
| |
| output = self._prepare_output_sequence( |
| model, |
| states_for_all_steps, |
| ) |
| return output, states_for_all_steps |
| |
| def apply(self, model, input_t, seq_lengths, states, timestep): |
| input_t = self.prepare_input(model, input_t) |
| states = self._apply( |
| model, input_t, seq_lengths, states, timestep) |
| output = self._prepare_output(model, states) |
| return output, states |
| |
| def _apply( |
| self, |
| model, input_t, seq_lengths, states, timestep, extra_inputs=None |
| ): |
| ''' |
| This method uses apply_override provided by a custom cell. |
| On the top it takes care of applying self.scope() to all the outputs. |
| While all the inputs stay within the scope this function was called |
| from. |
| ''' |
| args = self._rectify_apply_inputs( |
| input_t, seq_lengths, states, timestep, extra_inputs) |
| with core.NameScope(self.name): |
| return self.apply_override(model, *args) |
| |
| def _rectify_apply_inputs( |
| self, input_t, seq_lengths, states, timestep, extra_inputs): |
| ''' |
| Before applying a scope we make sure that all external blob names |
| are converted to blob reference. So further scoping doesn't affect them |
| ''' |
| |
| input_t, seq_lengths, timestep = _RectifyNames( |
| [input_t, seq_lengths, timestep]) |
| states = _RectifyNames(states) |
| if extra_inputs: |
| extra_input_names, extra_input_sizes = zip(*extra_inputs) |
| extra_inputs = _RectifyNames(extra_input_names) |
| extra_inputs = zip(extra_input_names, extra_input_sizes) |
| |
| arg_names = inspect.getargspec(self.apply_override).args |
| rectified = [input_t, seq_lengths, states, timestep] |
| if 'extra_inputs' in arg_names: |
| rectified.append(extra_inputs) |
| return rectified |
| |
| |
| def apply_override( |
| self, |
| model, input_t, seq_lengths, timestep, extra_inputs=None, |
| ): |
| ''' |
| A single step of a recurrent network to be implemented by each custom |
| RNNCell. |
| |
| model: ModelHelper object new operators would be added to |
| |
| input_t: singlse input with shape (1, batch_size, input_dim) |
| |
| seq_lengths: blob containing sequence lengths which would be passed to |
| LSTMUnit operator |
| |
| states: previous recurrent states |
| |
| timestep: current recurrent iteration. Could be used together with |
| seq_lengths in order to determine, if some shorter sequences |
| in the batch have already ended. |
| |
| extra_inputs: list of tuples (input, dim). specifies additional input |
| which is not subject to prepare_input(). (useful when a cell is a |
| component of a larger recurrent structure, e.g., attention) |
| ''' |
| raise NotImplementedError('Abstract method') |
| |
| def prepare_input(self, model, input_blob): |
| ''' |
| If some operations in _apply method depend only on the input, |
| not on recurrent states, they could be computed in advance. |
| |
| model: ModelHelper object new operators would be added to |
| |
| input_blob: either the whole input sequence with shape |
| (sequence_length, batch_size, input_dim) or a single input with shape |
| (1, batch_size, input_dim). |
| ''' |
| return input_blob |
| |
| def get_output_state_index(self): |
| ''' |
| Return index into state list of the "primary" step-wise output. |
| ''' |
| return 0 |
| |
| def get_state_names(self): |
| ''' |
| Returns recurrent state names with self.name scoping applied |
| ''' |
| return [self.scope(name) for name in self.get_state_names_override()] |
| |
| def get_state_names_override(self): |
| ''' |
| Override this function in your custom cell. |
| It should return the names of the recurrent states. |
| |
| It's required by apply_over_sequence method in order to allocate |
| recurrent states for all steps with meaningful names. |
| ''' |
| raise NotImplementedError('Abstract method') |
| |
| def get_output_dim(self): |
| ''' |
| Specifies the dimension (number of units) of stepwise output. |
| ''' |
| raise NotImplementedError('Abstract method') |
| |
| def _prepare_output(self, model, states): |
| ''' |
| Allows arbitrary post-processing of primary output. |
| ''' |
| return states[self.get_output_state_index()] |
| |
| def _prepare_output_sequence(self, model, state_outputs): |
| ''' |
| Allows arbitrary post-processing of primary sequence output. |
| |
| (Note that state_outputs alternates between full-sequence and final |
| output for each state, thus the index multiplier 2.) |
| ''' |
| output_sequence_index = 2 * self.get_output_state_index() |
| return state_outputs[output_sequence_index] |
| |
| |
| class LSTMInitializer(object): |
| def __init__(self, hidden_size): |
| self.hidden_size = hidden_size |
| |
| def create_states(self, model): |
| return [ |
| model.create_param( |
| param_name='initial_hidden_state', |
| initializer=Initializer(operator_name='ConstantFill', |
| value=0.0), |
| shape=[self.hidden_size], |
| ), |
| model.create_param( |
| param_name='initial_cell_state', |
| initializer=Initializer(operator_name='ConstantFill', |
| value=0.0), |
| shape=[self.hidden_size], |
| ) |
| ] |
| |
| |
| # based on https://pytorch.org/docs/master/nn.html#torch.nn.RNNCell |
| class BasicRNNCell(RNNCell): |
| def __init__( |
| self, |
| input_size, |
| hidden_size, |
| forget_bias, |
| memory_optimization, |
| drop_states=False, |
| initializer=None, |
| activation=None, |
| **kwargs |
| ): |
| super(BasicRNNCell, self).__init__(**kwargs) |
| self.drop_states = drop_states |
| self.input_size = input_size |
| self.hidden_size = hidden_size |
| self.activation = activation |
| |
| if self.activation not in ['relu', 'tanh']: |
| raise RuntimeError( |
| 'BasicRNNCell with unknown activation function (%s)' |
| % self.activation) |
| |
| def apply_override( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| hidden_t_prev = states[0] |
| |
| gates_t = brew.fc( |
| model, |
| hidden_t_prev, |
| 'gates_t', |
| dim_in=self.hidden_size, |
| dim_out=self.hidden_size, |
| axis=2, |
| ) |
| |
| brew.sum(model, [gates_t, input_t], gates_t) |
| if self.activation == 'tanh': |
| hidden_t = model.net.Tanh(gates_t, 'hidden_t') |
| elif self.activation == 'relu': |
| hidden_t = model.net.Relu(gates_t, 'hidden_t') |
| else: |
| raise RuntimeError( |
| 'BasicRNNCell with unknown activation function (%s)' |
| % self.activation) |
| |
| if seq_lengths is not None: |
| # TODO If this codepath becomes popular, it may be worth |
| # taking a look at optimizing it - for now a simple |
| # implementation is used to round out compatibility with |
| # ONNX. |
| timestep = model.net.CopyFromCPUInput( |
| timestep, 'timestep_gpu') |
| valid_b = model.net.GT( |
| [seq_lengths, timestep], 'valid_b', broadcast=1) |
| invalid_b = model.net.LE( |
| [seq_lengths, timestep], 'invalid_b', broadcast=1) |
| valid = model.net.Cast(valid_b, 'valid', to='float') |
| invalid = model.net.Cast(invalid_b, 'invalid', to='float') |
| |
| hidden_valid = model.net.Mul( |
| [hidden_t, valid], |
| 'hidden_valid', |
| broadcast=1, |
| axis=1, |
| ) |
| if self.drop_states: |
| hidden_t = hidden_valid |
| else: |
| hidden_invalid = model.net.Mul( |
| [hidden_t_prev, invalid], |
| 'hidden_invalid', |
| broadcast=1, axis=1) |
| hidden_t = model.net.Add( |
| [hidden_valid, hidden_invalid], hidden_t) |
| return (hidden_t,) |
| |
| def prepare_input(self, model, input_blob): |
| return brew.fc( |
| model, |
| input_blob, |
| self.scope('i2h'), |
| dim_in=self.input_size, |
| dim_out=self.hidden_size, |
| axis=2, |
| ) |
| |
| def get_state_names(self): |
| return (self.scope('hidden_t'),) |
| |
| def get_output_dim(self): |
| return self.hidden_size |
| |
| |
| class LSTMCell(RNNCell): |
| |
| def __init__( |
| self, |
| input_size, |
| hidden_size, |
| forget_bias, |
| memory_optimization, |
| drop_states=False, |
| initializer=None, |
| **kwargs |
| ): |
| super(LSTMCell, self).__init__(initializer=initializer, **kwargs) |
| self.initializer = initializer or LSTMInitializer( |
| hidden_size=hidden_size) |
| |
| self.input_size = input_size |
| self.hidden_size = hidden_size |
| self.forget_bias = float(forget_bias) |
| self.memory_optimization = memory_optimization |
| self.drop_states = drop_states |
| self.gates_size = 4 * self.hidden_size |
| |
| def apply_override( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| hidden_t_prev, cell_t_prev = states |
| |
| fc_input = hidden_t_prev |
| fc_input_dim = self.hidden_size |
| |
| if extra_inputs is not None: |
| extra_input_blobs, extra_input_sizes = zip(*extra_inputs) |
| fc_input = brew.concat( |
| model, |
| [hidden_t_prev] + list(extra_input_blobs), |
| 'gates_concatenated_input_t', |
| axis=2, |
| ) |
| fc_input_dim += sum(extra_input_sizes) |
| |
| gates_t = brew.fc( |
| model, |
| fc_input, |
| 'gates_t', |
| dim_in=fc_input_dim, |
| dim_out=self.gates_size, |
| axis=2, |
| ) |
| brew.sum(model, [gates_t, input_t], gates_t) |
| |
| if seq_lengths is not None: |
| inputs = [hidden_t_prev, cell_t_prev, gates_t, seq_lengths, timestep] |
| else: |
| inputs = [hidden_t_prev, cell_t_prev, gates_t, timestep] |
| |
| hidden_t, cell_t = model.net.LSTMUnit( |
| inputs, |
| ['hidden_state', 'cell_state'], |
| forget_bias=self.forget_bias, |
| drop_states=self.drop_states, |
| sequence_lengths=(seq_lengths is not None), |
| ) |
| model.net.AddExternalOutputs(hidden_t, cell_t) |
| if self.memory_optimization: |
| self.recompute_blobs = [gates_t] |
| |
| return hidden_t, cell_t |
| |
| def get_input_params(self): |
| return { |
| 'weights': self.scope('i2h') + '_w', |
| 'biases': self.scope('i2h') + '_b', |
| } |
| |
| def get_recurrent_params(self): |
| return { |
| 'weights': self.scope('gates_t') + '_w', |
| 'biases': self.scope('gates_t') + '_b', |
| } |
| |
| def prepare_input(self, model, input_blob): |
| return brew.fc( |
| model, |
| input_blob, |
| self.scope('i2h'), |
| dim_in=self.input_size, |
| dim_out=self.gates_size, |
| axis=2, |
| ) |
| |
| def get_state_names_override(self): |
| return ['hidden_t', 'cell_t'] |
| |
| def get_output_dim(self): |
| return self.hidden_size |
| |
| |
| class LayerNormLSTMCell(RNNCell): |
| |
| def __init__( |
| self, |
| input_size, |
| hidden_size, |
| forget_bias, |
| memory_optimization, |
| drop_states=False, |
| initializer=None, |
| **kwargs |
| ): |
| super(LayerNormLSTMCell, self).__init__( |
| initializer=initializer, **kwargs |
| ) |
| self.initializer = initializer or LSTMInitializer( |
| hidden_size=hidden_size |
| ) |
| |
| self.input_size = input_size |
| self.hidden_size = hidden_size |
| self.forget_bias = float(forget_bias) |
| self.memory_optimization = memory_optimization |
| self.drop_states = drop_states |
| self.gates_size = 4 * self.hidden_size |
| |
| def _apply( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| hidden_t_prev, cell_t_prev = states |
| |
| fc_input = hidden_t_prev |
| fc_input_dim = self.hidden_size |
| |
| if extra_inputs is not None: |
| extra_input_blobs, extra_input_sizes = zip(*extra_inputs) |
| fc_input = brew.concat( |
| model, |
| [hidden_t_prev] + list(extra_input_blobs), |
| self.scope('gates_concatenated_input_t'), |
| axis=2, |
| ) |
| fc_input_dim += sum(extra_input_sizes) |
| |
| gates_t = brew.fc( |
| model, |
| fc_input, |
| self.scope('gates_t'), |
| dim_in=fc_input_dim, |
| dim_out=self.gates_size, |
| axis=2, |
| ) |
| brew.sum(model, [gates_t, input_t], gates_t) |
| |
| # brew.layer_norm call is only difference from LSTMCell |
| gates_t, _, _ = brew.layer_norm( |
| model, |
| self.scope('gates_t'), |
| self.scope('gates_t_norm'), |
| dim_in=self.gates_size, |
| axis=-1, |
| ) |
| |
| hidden_t, cell_t = model.net.LSTMUnit( |
| [ |
| hidden_t_prev, |
| cell_t_prev, |
| gates_t, |
| seq_lengths, |
| timestep, |
| ], |
| self.get_state_names(), |
| forget_bias=self.forget_bias, |
| drop_states=self.drop_states, |
| ) |
| model.net.AddExternalOutputs(hidden_t, cell_t) |
| if self.memory_optimization: |
| self.recompute_blobs = [gates_t] |
| |
| return hidden_t, cell_t |
| |
| def get_input_params(self): |
| return { |
| 'weights': self.scope('i2h') + '_w', |
| 'biases': self.scope('i2h') + '_b', |
| } |
| |
| def prepare_input(self, model, input_blob): |
| return brew.fc( |
| model, |
| input_blob, |
| self.scope('i2h'), |
| dim_in=self.input_size, |
| dim_out=self.gates_size, |
| axis=2, |
| ) |
| |
| def get_state_names(self): |
| return (self.scope('hidden_t'), self.scope('cell_t')) |
| |
| |
| class MILSTMCell(LSTMCell): |
| |
| def _apply( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| hidden_t_prev, cell_t_prev = states |
| |
| fc_input = hidden_t_prev |
| fc_input_dim = self.hidden_size |
| |
| if extra_inputs is not None: |
| extra_input_blobs, extra_input_sizes = zip(*extra_inputs) |
| fc_input = brew.concat( |
| model, |
| [hidden_t_prev] + list(extra_input_blobs), |
| self.scope('gates_concatenated_input_t'), |
| axis=2, |
| ) |
| fc_input_dim += sum(extra_input_sizes) |
| |
| prev_t = brew.fc( |
| model, |
| fc_input, |
| self.scope('prev_t'), |
| dim_in=fc_input_dim, |
| dim_out=self.gates_size, |
| axis=2, |
| ) |
| |
| # defining initializers for MI parameters |
| alpha = model.create_param( |
| self.scope('alpha'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=1.0), |
| ) |
| beta_h = model.create_param( |
| self.scope('beta1'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=1.0), |
| ) |
| beta_i = model.create_param( |
| self.scope('beta2'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=1.0), |
| ) |
| b = model.create_param( |
| self.scope('b'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=0.0), |
| ) |
| |
| # alpha * input_t + beta_h |
| # Shape: [1, batch_size, 4 * hidden_size] |
| alpha_by_input_t_plus_beta_h = model.net.ElementwiseLinear( |
| [input_t, alpha, beta_h], |
| self.scope('alpha_by_input_t_plus_beta_h'), |
| axis=2, |
| ) |
| # (alpha * input_t + beta_h) * prev_t = |
| # alpha * input_t * prev_t + beta_h * prev_t |
| # Shape: [1, batch_size, 4 * hidden_size] |
| alpha_by_input_t_plus_beta_h_by_prev_t = model.net.Mul( |
| [alpha_by_input_t_plus_beta_h, prev_t], |
| self.scope('alpha_by_input_t_plus_beta_h_by_prev_t') |
| ) |
| # beta_i * input_t + b |
| # Shape: [1, batch_size, 4 * hidden_size] |
| beta_i_by_input_t_plus_b = model.net.ElementwiseLinear( |
| [input_t, beta_i, b], |
| self.scope('beta_i_by_input_t_plus_b'), |
| axis=2, |
| ) |
| # alpha * input_t * prev_t + beta_h * prev_t + beta_i * input_t + b |
| # Shape: [1, batch_size, 4 * hidden_size] |
| gates_t = brew.sum( |
| model, |
| [alpha_by_input_t_plus_beta_h_by_prev_t, beta_i_by_input_t_plus_b], |
| self.scope('gates_t') |
| ) |
| hidden_t, cell_t = model.net.LSTMUnit( |
| [hidden_t_prev, cell_t_prev, gates_t, seq_lengths, timestep], |
| [self.scope('hidden_t_intermediate'), self.scope('cell_t')], |
| forget_bias=self.forget_bias, |
| drop_states=self.drop_states, |
| ) |
| model.net.AddExternalOutputs( |
| cell_t, |
| hidden_t, |
| ) |
| if self.memory_optimization: |
| self.recompute_blobs = [gates_t] |
| return hidden_t, cell_t |
| |
| |
| class LayerNormMILSTMCell(LSTMCell): |
| |
| def _apply( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| hidden_t_prev, cell_t_prev = states |
| |
| fc_input = hidden_t_prev |
| fc_input_dim = self.hidden_size |
| |
| if extra_inputs is not None: |
| extra_input_blobs, extra_input_sizes = zip(*extra_inputs) |
| fc_input = brew.concat( |
| model, |
| [hidden_t_prev] + list(extra_input_blobs), |
| self.scope('gates_concatenated_input_t'), |
| axis=2, |
| ) |
| fc_input_dim += sum(extra_input_sizes) |
| |
| prev_t = brew.fc( |
| model, |
| fc_input, |
| self.scope('prev_t'), |
| dim_in=fc_input_dim, |
| dim_out=self.gates_size, |
| axis=2, |
| ) |
| |
| # defining initializers for MI parameters |
| alpha = model.create_param( |
| self.scope('alpha'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=1.0), |
| ) |
| beta_h = model.create_param( |
| self.scope('beta1'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=1.0), |
| ) |
| beta_i = model.create_param( |
| self.scope('beta2'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=1.0), |
| ) |
| b = model.create_param( |
| self.scope('b'), |
| shape=[self.gates_size], |
| initializer=Initializer('ConstantFill', value=0.0), |
| ) |
| |
| # alpha * input_t + beta_h |
| # Shape: [1, batch_size, 4 * hidden_size] |
| alpha_by_input_t_plus_beta_h = model.net.ElementwiseLinear( |
| [input_t, alpha, beta_h], |
| self.scope('alpha_by_input_t_plus_beta_h'), |
| axis=2, |
| ) |
| # (alpha * input_t + beta_h) * prev_t = |
| # alpha * input_t * prev_t + beta_h * prev_t |
| # Shape: [1, batch_size, 4 * hidden_size] |
| alpha_by_input_t_plus_beta_h_by_prev_t = model.net.Mul( |
| [alpha_by_input_t_plus_beta_h, prev_t], |
| self.scope('alpha_by_input_t_plus_beta_h_by_prev_t') |
| ) |
| # beta_i * input_t + b |
| # Shape: [1, batch_size, 4 * hidden_size] |
| beta_i_by_input_t_plus_b = model.net.ElementwiseLinear( |
| [input_t, beta_i, b], |
| self.scope('beta_i_by_input_t_plus_b'), |
| axis=2, |
| ) |
| # alpha * input_t * prev_t + beta_h * prev_t + beta_i * input_t + b |
| # Shape: [1, batch_size, 4 * hidden_size] |
| gates_t = brew.sum( |
| model, |
| [alpha_by_input_t_plus_beta_h_by_prev_t, beta_i_by_input_t_plus_b], |
| self.scope('gates_t') |
| ) |
| # brew.layer_norm call is only difference from MILSTMCell._apply |
| gates_t, _, _ = brew.layer_norm( |
| model, |
| self.scope('gates_t'), |
| self.scope('gates_t_norm'), |
| dim_in=self.gates_size, |
| axis=-1, |
| ) |
| hidden_t, cell_t = model.net.LSTMUnit( |
| [hidden_t_prev, cell_t_prev, gates_t, seq_lengths, timestep], |
| [self.scope('hidden_t_intermediate'), self.scope('cell_t')], |
| forget_bias=self.forget_bias, |
| drop_states=self.drop_states, |
| ) |
| model.net.AddExternalOutputs( |
| cell_t, |
| hidden_t, |
| ) |
| if self.memory_optimization: |
| self.recompute_blobs = [gates_t] |
| return hidden_t, cell_t |
| |
| |
| class DropoutCell(RNNCell): |
| ''' |
| Wraps arbitrary RNNCell, applying dropout to its output (but not to the |
| recurrent connection for the corresponding state). |
| ''' |
| |
| def __init__( |
| self, |
| internal_cell, |
| dropout_ratio=None, |
| use_cudnn=False, |
| **kwargs |
| ): |
| self.internal_cell = internal_cell |
| self.dropout_ratio = dropout_ratio |
| assert 'is_test' in kwargs, "Argument 'is_test' is required" |
| self.is_test = kwargs.pop('is_test') |
| self.use_cudnn = use_cudnn |
| super(DropoutCell, self).__init__(**kwargs) |
| |
| self.prepare_input = internal_cell.prepare_input |
| self.get_output_state_index = internal_cell.get_output_state_index |
| self.get_state_names = internal_cell.get_state_names |
| self.get_output_dim = internal_cell.get_output_dim |
| |
| self.mask = 0 |
| |
| def _apply( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| return self.internal_cell._apply( |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs, |
| ) |
| |
| def _prepare_output(self, model, states): |
| output = self.internal_cell._prepare_output( |
| model, |
| states, |
| ) |
| if self.dropout_ratio is not None: |
| output = self._apply_dropout(model, output) |
| return output |
| |
| def _prepare_output_sequence(self, model, state_outputs): |
| output = self.internal_cell._prepare_output_sequence( |
| model, |
| state_outputs, |
| ) |
| if self.dropout_ratio is not None: |
| output = self._apply_dropout(model, output) |
| return output |
| |
| def _apply_dropout(self, model, output): |
| if self.dropout_ratio and not self.forward_only: |
| with core.NameScope(self.name or ''): |
| output = brew.dropout( |
| model, |
| output, |
| str(output) + '_with_dropout_mask{}'.format(self.mask), |
| ratio=float(self.dropout_ratio), |
| is_test=self.is_test, |
| use_cudnn=self.use_cudnn, |
| ) |
| self.mask += 1 |
| return output |
| |
| |
| class MultiRNNCellInitializer(object): |
| def __init__(self, cells): |
| self.cells = cells |
| |
| def create_states(self, model): |
| states = [] |
| for i, cell in enumerate(self.cells): |
| if cell.initializer is None: |
| raise Exception("Either initial states " |
| "or initializer have to be set") |
| |
| with core.NameScope("layer_{}".format(i)),\ |
| core.NameScope(cell.name): |
| states.extend(cell.initializer.create_states(model)) |
| return states |
| |
| |
| class MultiRNNCell(RNNCell): |
| ''' |
| Multilayer RNN via the composition of RNNCell instance. |
| |
| It is the responsibility of calling code to ensure the compatibility |
| of the successive layers in terms of input/output dimensiality, etc., |
| and to ensure that their blobs do not have name conflicts, typically by |
| creating the cells with names that specify layer number. |
| |
| Assumes first state (recurrent output) for each layer should be the input |
| to the next layer. |
| ''' |
| |
| def __init__(self, cells, residual_output_layers=None, **kwargs): |
| ''' |
| cells: list of RNNCell instances, from input to output side. |
| |
| name: string designating network component (for scoping) |
| |
| residual_output_layers: list of indices of layers whose input will |
| be added elementwise to their output elementwise. (It is the |
| responsibility of the client code to ensure shape compatibility.) |
| Note that layer 0 (zero) cannot have residual output because of the |
| timing of prepare_input(). |
| |
| forward_only: used to construct inference-only network. |
| ''' |
| super(MultiRNNCell, self).__init__(**kwargs) |
| self.cells = cells |
| |
| if residual_output_layers is None: |
| self.residual_output_layers = [] |
| else: |
| self.residual_output_layers = residual_output_layers |
| |
| output_index_per_layer = [] |
| base_index = 0 |
| for cell in self.cells: |
| output_index_per_layer.append( |
| base_index + cell.get_output_state_index(), |
| ) |
| base_index += len(cell.get_state_names()) |
| |
| self.output_connected_layers = [] |
| self.output_indices = [] |
| for i in range(len(self.cells) - 1): |
| if (i + 1) in self.residual_output_layers: |
| self.output_connected_layers.append(i) |
| self.output_indices.append(output_index_per_layer[i]) |
| else: |
| self.output_connected_layers = [] |
| self.output_indices = [] |
| self.output_connected_layers.append(len(self.cells) - 1) |
| self.output_indices.append(output_index_per_layer[-1]) |
| |
| self.state_names = [] |
| for i, cell in enumerate(self.cells): |
| self.state_names.extend( |
| map(self.layer_scoper(i), cell.get_state_names()) |
| ) |
| |
| self.initializer = MultiRNNCellInitializer(cells) |
| |
| def layer_scoper(self, layer_id): |
| def helper(name): |
| return "{}/layer_{}/{}".format(self.name, layer_id, name) |
| return helper |
| |
| def prepare_input(self, model, input_blob): |
| input_blob = _RectifyName(input_blob) |
| with core.NameScope(self.name or ''): |
| return self.cells[0].prepare_input(model, input_blob) |
| |
| def _apply( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| ''' |
| Because below we will do scoping across layers, we need |
| to make sure that string blob names are convereted to BlobReference |
| objects. |
| ''' |
| |
| input_t, seq_lengths, states, timestep, extra_inputs = \ |
| self._rectify_apply_inputs( |
| input_t, seq_lengths, states, timestep, extra_inputs) |
| |
| states_per_layer = [len(cell.get_state_names()) for cell in self.cells] |
| assert len(states) == sum(states_per_layer) |
| |
| next_states = [] |
| states_index = 0 |
| |
| layer_input = input_t |
| for i, layer_cell in enumerate(self.cells): |
| # # If cells don't have different names we still |
| # take care of scoping |
| with core.NameScope(self.name), core.NameScope("layer_{}".format(i)): |
| num_states = states_per_layer[i] |
| layer_states = states[states_index:(states_index + num_states)] |
| states_index += num_states |
| |
| if i > 0: |
| prepared_input = layer_cell.prepare_input( |
| model, layer_input) |
| else: |
| prepared_input = layer_input |
| |
| layer_next_states = layer_cell._apply( |
| model, |
| prepared_input, |
| seq_lengths, |
| layer_states, |
| timestep, |
| extra_inputs=(None if i > 0 else extra_inputs), |
| ) |
| # Since we're using here non-public method _apply, |
| # instead of apply, we have to manually extract output |
| # from states |
| if i != len(self.cells) - 1: |
| layer_output = layer_cell._prepare_output( |
| model, |
| layer_next_states, |
| ) |
| if i > 0 and i in self.residual_output_layers: |
| layer_input = brew.sum( |
| model, |
| [layer_output, layer_input], |
| self.scope('residual_output_{}'.format(i)), |
| ) |
| else: |
| layer_input = layer_output |
| |
| next_states.extend(layer_next_states) |
| return next_states |
| |
| def get_state_names(self): |
| return self.state_names |
| |
| def get_output_state_index(self): |
| index = 0 |
| for cell in self.cells[:-1]: |
| index += len(cell.get_state_names()) |
| index += self.cells[-1].get_output_state_index() |
| return index |
| |
| def _prepare_output(self, model, states): |
| connected_outputs = [] |
| state_index = 0 |
| for i, cell in enumerate(self.cells): |
| num_states = len(cell.get_state_names()) |
| if i in self.output_connected_layers: |
| layer_states = states[state_index:state_index + num_states] |
| layer_output = cell._prepare_output( |
| model, |
| layer_states |
| ) |
| connected_outputs.append(layer_output) |
| state_index += num_states |
| if len(connected_outputs) > 1: |
| output = brew.sum( |
| model, |
| connected_outputs, |
| self.scope('residual_output'), |
| ) |
| else: |
| output = connected_outputs[0] |
| return output |
| |
| def _prepare_output_sequence(self, model, states): |
| connected_outputs = [] |
| state_index = 0 |
| for i, cell in enumerate(self.cells): |
| num_states = 2 * len(cell.get_state_names()) |
| if i in self.output_connected_layers: |
| layer_states = states[state_index:state_index + num_states] |
| layer_output = cell._prepare_output_sequence( |
| model, |
| layer_states |
| ) |
| connected_outputs.append(layer_output) |
| state_index += num_states |
| if len(connected_outputs) > 1: |
| output = brew.sum( |
| model, |
| connected_outputs, |
| self.scope('residual_output_sequence'), |
| ) |
| else: |
| output = connected_outputs[0] |
| return output |
| |
| |
| class AttentionCell(RNNCell): |
| |
| def __init__( |
| self, |
| encoder_output_dim, |
| encoder_outputs, |
| encoder_lengths, |
| decoder_cell, |
| decoder_state_dim, |
| attention_type, |
| weighted_encoder_outputs, |
| attention_memory_optimization, |
| **kwargs |
| ): |
| super(AttentionCell, self).__init__(**kwargs) |
| self.encoder_output_dim = encoder_output_dim |
| self.encoder_outputs = encoder_outputs |
| self.encoder_lengths = encoder_lengths |
| self.decoder_cell = decoder_cell |
| self.decoder_state_dim = decoder_state_dim |
| self.weighted_encoder_outputs = weighted_encoder_outputs |
| self.encoder_outputs_transposed = None |
| assert attention_type in [ |
| AttentionType.Regular, |
| AttentionType.Recurrent, |
| AttentionType.Dot, |
| AttentionType.SoftCoverage, |
| ] |
| self.attention_type = attention_type |
| self.attention_memory_optimization = attention_memory_optimization |
| |
| def _apply( |
| self, |
| model, |
| input_t, |
| seq_lengths, |
| states, |
| timestep, |
| extra_inputs=None, |
| ): |
| if self.attention_type == AttentionType.SoftCoverage: |
| decoder_prev_states = states[:-2] |
| attention_weighted_encoder_context_t_prev = states[-2] |
| coverage_t_prev = states[-1] |
| else: |
| decoder_prev_states = states[:-1] |
| attention_weighted_encoder_context_t_prev = states[-1] |
| |
| assert extra_inputs is None |
| |
| decoder_states = self.decoder_cell._apply( |
| model, |
| input_t, |
| seq_lengths, |
| decoder_prev_states, |
| timestep, |
| extra_inputs=[( |
| attention_weighted_encoder_context_t_prev, |
| self.encoder_output_dim, |
| )], |
| ) |
| |
| self.hidden_t_intermediate = self.decoder_cell._prepare_output( |
| model, |
| decoder_states, |
| ) |
| |
| if self.attention_type == AttentionType.Recurrent: |
| ( |
| attention_weighted_encoder_context_t, |
| self.attention_weights_3d, |
| attention_blobs, |
| ) = apply_recurrent_attention( |
| model=model, |
| encoder_output_dim=self.encoder_output_dim, |
| encoder_outputs_transposed=self.encoder_outputs_transposed, |
| weighted_encoder_outputs=self.weighted_encoder_outputs, |
| decoder_hidden_state_t=self.hidden_t_intermediate, |
| decoder_hidden_state_dim=self.decoder_state_dim, |
| scope=self.name, |
| attention_weighted_encoder_context_t_prev=( |
| attention_weighted_encoder_context_t_prev |
| ), |
| encoder_lengths=self.encoder_lengths, |
| ) |
| elif self.attention_type == AttentionType.Regular: |
| ( |
| attention_weighted_encoder_context_t, |
| self.attention_weights_3d, |
| attention_blobs, |
| ) = apply_regular_attention( |
| model=model, |
| encoder_output_dim=self.encoder_output_dim, |
| encoder_outputs_transposed=self.encoder_outputs_transposed, |
| weighted_encoder_outputs=self.weighted_encoder_outputs, |
| decoder_hidden_state_t=self.hidden_t_intermediate, |
| decoder_hidden_state_dim=self.decoder_state_dim, |
| scope=self.name, |
| encoder_lengths=self.encoder_lengths, |
| ) |
| elif self.attention_type == AttentionType.Dot: |
| ( |
| attention_weighted_encoder_context_t, |
| self.attention_weights_3d, |
| attention_blobs, |
| ) = apply_dot_attention( |
| model=model, |
| encoder_output_dim=self.encoder_output_dim, |
| encoder_outputs_transposed=self.encoder_outputs_transposed, |
| decoder_hidden_state_t=self.hidden_t_intermediate, |
| decoder_hidden_state_dim=self.decoder_state_dim, |
| scope=self.name, |
| encoder_lengths=self.encoder_lengths, |
| ) |
| elif self.attention_type == AttentionType.SoftCoverage: |
| ( |
| attention_weighted_encoder_context_t, |
| self.attention_weights_3d, |
| attention_blobs, |
| coverage_t, |
| ) = apply_soft_coverage_attention( |
| model=model, |
| encoder_output_dim=self.encoder_output_dim, |
| encoder_outputs_transposed=self.encoder_outputs_transposed, |
| weighted_encoder_outputs=self.weighted_encoder_outputs, |
| decoder_hidden_state_t=self.hidden_t_intermediate, |
| decoder_hidden_state_dim=self.decoder_state_dim, |
| scope=self.name, |
| encoder_lengths=self.encoder_lengths, |
| coverage_t_prev=coverage_t_prev, |
| coverage_weights=self.coverage_weights, |
| ) |
| else: |
| raise Exception('Attention type {} not implemented'.format( |
| self.attention_type |
| )) |
| |
| if self.attention_memory_optimization: |
| self.recompute_blobs.extend(attention_blobs) |
| |
| output = list(decoder_states) + [attention_weighted_encoder_context_t] |
| if self.attention_type == AttentionType.SoftCoverage: |
| output.append(coverage_t) |
| |
| output[self.decoder_cell.get_output_state_index()] = model.Copy( |
| output[self.decoder_cell.get_output_state_index()], |
| self.scope('hidden_t_external'), |
| ) |
| model.net.AddExternalOutputs(*output) |
| |
| return output |
| |
| def get_attention_weights(self): |
| # [batch_size, encoder_length, 1] |
| return self.attention_weights_3d |
| |
| def prepare_input(self, model, input_blob): |
| if self.encoder_outputs_transposed is None: |
| self.encoder_outputs_transposed = brew.transpose( |
| model, |
| self.encoder_outputs, |
| self.scope('encoder_outputs_transposed'), |
| axes=[1, 2, 0], |
| ) |
| if ( |
| self.weighted_encoder_outputs is None and |
| self.attention_type != AttentionType.Dot |
| ): |
| self.weighted_encoder_outputs = brew.fc( |
| model, |
| self.encoder_outputs, |
| self.scope('weighted_encoder_outputs'), |
| dim_in=self.encoder_output_dim, |
| dim_out=self.encoder_output_dim, |
| axis=2, |
| ) |
| |
| return self.decoder_cell.prepare_input(model, input_blob) |
| |
| def build_initial_coverage(self, model): |
| """ |
| initial_coverage is always zeros of shape [encoder_length], |
| which shape must be determined programmatically dureing network |
| computation. |
| |
| This method also sets self.coverage_weights, a separate transform |
| of encoder_outputs which is used to determine coverage contribution |
| tp attention. |
| """ |
| assert self.attention_type == AttentionType.SoftCoverage |
| |
| # [encoder_length, batch_size, encoder_output_dim] |
| self.coverage_weights = brew.fc( |
| model, |
| self.encoder_outputs, |
| self.scope('coverage_weights'), |
| dim_in=self.encoder_output_dim, |
| dim_out=self.encoder_output_dim, |
| axis=2, |
| ) |
| |
| encoder_length = model.net.Slice( |
| model.net.Shape(self.encoder_outputs), |
| starts=[0], |
| ends=[1], |
| ) |
| if ( |
| scope.CurrentDeviceScope() is not None and |
| core.IsGPUDeviceType(scope.CurrentDeviceScope().device_type) |
| ): |
| encoder_length = model.net.CopyGPUToCPU( |
| encoder_length, |
| 'encoder_length_cpu', |
| ) |
| # total attention weight applied across decoding steps_per_checkpoint |
| # shape: [encoder_length] |
| initial_coverage = model.net.ConstantFill( |
| encoder_length, |
| self.scope('initial_coverage'), |
| value=0.0, |
| input_as_shape=1, |
| ) |
| return initial_coverage |
| |
| def get_state_names(self): |
| state_names = list(self.decoder_cell.get_state_names()) |
| state_names[self.get_output_state_index()] = self.scope( |
| 'hidden_t_external', |
| ) |
| state_names.append(self.scope('attention_weighted_encoder_context_t')) |
| if self.attention_type == AttentionType.SoftCoverage: |
| state_names.append(self.scope('coverage_t')) |
| return state_names |
| |
| def get_output_dim(self): |
| return self.decoder_state_dim + self.encoder_output_dim |
| |
| def get_output_state_index(self): |
| return self.decoder_cell.get_output_state_index() |
| |
| def _prepare_output(self, model, states): |
| if self.attention_type == AttentionType.SoftCoverage: |
| attention_context = states[-2] |
| else: |
| attention_context = states[-1] |
| |
| with core.NameScope(self.name or ''): |
| output = brew.concat( |
| model, |
| [self.hidden_t_intermediate, attention_context], |
| 'states_and_context_combination', |
| axis=2, |
| ) |
| |
| return output |
| |
| def _prepare_output_sequence(self, model, state_outputs): |
| if self.attention_type == AttentionType.SoftCoverage: |
| decoder_state_outputs = state_outputs[:-4] |
| else: |
| decoder_state_outputs = state_outputs[:-2] |
| |
| decoder_output = self.decoder_cell._prepare_output_sequence( |
| model, |
| decoder_state_outputs, |
| ) |
| |
| if self.attention_type == AttentionType.SoftCoverage: |
| attention_context_index = 2 * (len(self.get_state_names()) - 2) |
| else: |
| attention_context_index = 2 * (len(self.get_state_names()) - 1) |
| |
| with core.NameScope(self.name or ''): |
| output = brew.concat( |
| model, |
| [ |
| decoder_output, |
| state_outputs[attention_context_index], |
| ], |
| 'states_and_context_combination', |
| axis=2, |
| ) |
| return output |
| |
| |
| class LSTMWithAttentionCell(AttentionCell): |
| |
| def __init__( |
| self, |
| encoder_output_dim, |
| encoder_outputs, |
| encoder_lengths, |
| decoder_input_dim, |
| decoder_state_dim, |
| name, |
| attention_type, |
| weighted_encoder_outputs, |
| forget_bias, |
| lstm_memory_optimization, |
| attention_memory_optimization, |
| forward_only=False, |
| ): |
| decoder_cell = LSTMCell( |
| input_size=decoder_input_dim, |
| hidden_size=decoder_state_dim, |
| forget_bias=forget_bias, |
| memory_optimization=lstm_memory_optimization, |
| name='{}/decoder'.format(name), |
| forward_only=False, |
| drop_states=False, |
| ) |
| super(LSTMWithAttentionCell, self).__init__( |
| encoder_output_dim=encoder_output_dim, |
| encoder_outputs=encoder_outputs, |
| encoder_lengths=encoder_lengths, |
| decoder_cell=decoder_cell, |
| decoder_state_dim=decoder_state_dim, |
| name=name, |
| attention_type=attention_type, |
| weighted_encoder_outputs=weighted_encoder_outputs, |
| attention_memory_optimization=attention_memory_optimization, |
| forward_only=forward_only, |
| ) |
| |
| |
| class MILSTMWithAttentionCell(AttentionCell): |
| |
| def __init__( |
| self, |
| encoder_output_dim, |
| encoder_outputs, |
| decoder_input_dim, |
| decoder_state_dim, |
| name, |
| attention_type, |
| weighted_encoder_outputs, |
| forget_bias, |
| lstm_memory_optimization, |
| attention_memory_optimization, |
| forward_only=False, |
| ): |
| decoder_cell = MILSTMCell( |
| input_size=decoder_input_dim, |
| hidden_size=decoder_state_dim, |
| forget_bias=forget_bias, |
| memory_optimization=lstm_memory_optimization, |
| name='{}/decoder'.format(name), |
| forward_only=False, |
| drop_states=False, |
| ) |
| super(MILSTMWithAttentionCell, self).__init__( |
| encoder_output_dim=encoder_output_dim, |
| encoder_outputs=encoder_outputs, |
| decoder_cell=decoder_cell, |
| decoder_state_dim=decoder_state_dim, |
| name=name, |
| attention_type=attention_type, |
| weighted_encoder_outputs=weighted_encoder_outputs, |
| attention_memory_optimization=attention_memory_optimization, |
| forward_only=forward_only, |
| ) |
| |
| |
| def _LSTM( |
| cell_class, |
| model, |
| input_blob, |
| seq_lengths, |
| initial_states, |
| dim_in, |
| dim_out, |
| scope=None, |
| outputs_with_grads=(0,), |
| return_params=False, |
| memory_optimization=False, |
| forget_bias=0.0, |
| forward_only=False, |
| drop_states=False, |
| return_last_layer_only=True, |
| static_rnn_unroll_size=None, |
| **cell_kwargs |
| ): |
| ''' |
| Adds a standard LSTM recurrent network operator to a model. |
| |
| cell_class: LSTMCell or compatible subclass |
| |
| model: ModelHelper object new operators would be added to |
| |
| input_blob: the input sequence in a format T x N x D |
| where T is sequence size, N - batch size and D - input dimension |
| |
| seq_lengths: blob containing sequence lengths which would be passed to |
| LSTMUnit operator |
| |
| initial_states: a list of (2 * num_layers) blobs representing the initial |
| hidden and cell states of each layer. If this argument is None, |
| these states will be added to the model as network parameters. |
| |
| dim_in: input dimension |
| |
| dim_out: number of units per LSTM layer |
| (use int for single-layer LSTM, list of ints for multi-layer) |
| |
| outputs_with_grads : position indices of output blobs for LAST LAYER which |
| will receive external error gradient during backpropagation. |
| These outputs are: (h_all, h_last, c_all, c_last) |
| |
| return_params: if True, will return a dictionary of parameters of the LSTM |
| |
| memory_optimization: if enabled, the LSTM step is recomputed on backward |
| step so that we don't need to store forward activations for each |
| timestep. Saves memory with cost of computation. |
| |
| forget_bias: forget gate bias (default 0.0) |
| |
| forward_only: whether to create a backward pass |
| |
| drop_states: drop invalid states, passed through to LSTMUnit operator |
| |
| return_last_layer_only: only return outputs from final layer |
| (so that length of results does depend on number of layers) |
| |
| static_rnn_unroll_size: if not None, we will use static RNN which is |
| unrolled into Caffe2 graph. The size of the unroll is the value of |
| this parameter. |
| ''' |
| if type(dim_out) is not list and type(dim_out) is not tuple: |
| dim_out = [dim_out] |
| num_layers = len(dim_out) |
| |
| cells = [] |
| for i in range(num_layers): |
| cell = cell_class( |
| input_size=(dim_in if i == 0 else dim_out[i - 1]), |
| hidden_size=dim_out[i], |
| forget_bias=forget_bias, |
| memory_optimization=memory_optimization, |
| name=scope if num_layers == 1 else None, |
| forward_only=forward_only, |
| drop_states=drop_states, |
| **cell_kwargs |
| ) |
| cells.append(cell) |
| |
| cell = MultiRNNCell( |
| cells, |
| name=scope, |
| forward_only=forward_only, |
| ) if num_layers > 1 else cells[0] |
| |
| cell = ( |
| cell if static_rnn_unroll_size is None |
| else UnrolledCell(cell, static_rnn_unroll_size)) |
| |
| # outputs_with_grads argument indexes into final layer |
| outputs_with_grads = [4 * (num_layers - 1) + i for i in outputs_with_grads] |
| _, result = cell.apply_over_sequence( |
| model=model, |
| inputs=input_blob, |
| seq_lengths=seq_lengths, |
| initial_states=initial_states, |
| outputs_with_grads=outputs_with_grads, |
| ) |
| |
| if return_last_layer_only: |
| result = result[4 * (num_layers - 1):] |
| if return_params: |
| result = list(result) + [{ |
| 'input': cell.get_input_params(), |
| 'recurrent': cell.get_recurrent_params(), |
| }] |
| return tuple(result) |
| |
| |
| LSTM = functools.partial(_LSTM, LSTMCell) |
| BasicRNN = functools.partial(_LSTM, BasicRNNCell) |
| MILSTM = functools.partial(_LSTM, MILSTMCell) |
| LayerNormLSTM = functools.partial(_LSTM, LayerNormLSTMCell) |
| LayerNormMILSTM = functools.partial(_LSTM, LayerNormMILSTMCell) |
| |
| |
| class UnrolledCell(RNNCell): |
| def __init__(self, cell, T): |
| self.T = T |
| self.cell = cell |
| |
| def apply_over_sequence( |
| self, |
| model, |
| inputs, |
| seq_lengths, |
| initial_states, |
| outputs_with_grads=None, |
| ): |
| inputs = self.cell.prepare_input(model, inputs) |
| |
| # Now they are blob references - outputs of splitting the input sequence |
| split_inputs = model.net.Split( |
| inputs, |
| [str(inputs) + "_timestep_{}".format(i) |
| for i in range(self.T)], |
| axis=0) |
| if self.T == 1: |
| split_inputs = [split_inputs] |
| |
| states = initial_states |
| all_states = [] |
| for t in range(0, self.T): |
| scope_name = "timestep_{}".format(t) |
| # Parameters of all timesteps are shared |
| with ParameterSharing({scope_name: ''}),\ |
| scope.NameScope(scope_name): |
| timestep = model.param_init_net.ConstantFill( |
| [], "timestep", value=t, shape=[1], |
| dtype=core.DataType.INT32, |
| device_option=core.DeviceOption(caffe2_pb2.CPU)) |
| states = self.cell._apply( |
| model=model, |
| input_t=split_inputs[t], |
| seq_lengths=seq_lengths, |
| states=states, |
| timestep=timestep, |
| ) |
| all_states.append(states) |
| |
| all_states = zip(*all_states) |
| all_states = [ |
| model.net.Concat( |
| list(full_output), |
| [ |
| str(full_output[0])[len("timestep_0/"):] + "_concat", |
| str(full_output[0])[len("timestep_0/"):] + "_concat_info" |
| |
| ], |
| axis=0)[0] |
| for full_output in all_states |
| ] |
| # Interleave the state values similar to |
| # |
| # x = [1, 3, 5] |
| # y = [2, 4, 6] |
| # z = [val for pair in zip(x, y) for val in pair] |
| # # z is [1, 2, 3, 4, 5, 6] |
| # |
| # and returns it as outputs |
| outputs = tuple( |
| state for state_pair in zip(all_states, states) for state in state_pair |
| ) |
| outputs_without_grad = set(range(len(outputs))) - set( |
| outputs_with_grads) |
| for i in outputs_without_grad: |
| model.net.ZeroGradient(outputs[i], []) |
| logging.debug("Added 0 gradients for blobs:", |
| [outputs[i] for i in outputs_without_grad]) |
| |
| final_output = self.cell._prepare_output_sequence(model, outputs) |
| |
| return final_output, outputs |
| |
| |
| def GetLSTMParamNames(): |
| weight_params = ["input_gate_w", "forget_gate_w", "output_gate_w", "cell_w"] |
| bias_params = ["input_gate_b", "forget_gate_b", "output_gate_b", "cell_b"] |
| return {'weights': weight_params, 'biases': bias_params} |
| |
| |
| def InitFromLSTMParams(lstm_pblobs, param_values): |
| ''' |
| Set the parameters of LSTM based on predefined values |
| ''' |
| weight_params = GetLSTMParamNames()['weights'] |
| bias_params = GetLSTMParamNames()['biases'] |
| for input_type in viewkeys(param_values): |
| weight_values = [ |
| param_values[input_type][w].flatten() |
| for w in weight_params |
| ] |
| wmat = np.array([]) |
| for w in weight_values: |
| wmat = np.append(wmat, w) |
| bias_values = [ |
| param_values[input_type][b].flatten() |
| for b in bias_params |
| ] |
| bm = np.array([]) |
| for b in bias_values: |
| bm = np.append(bm, b) |
| |
| weights_blob = lstm_pblobs[input_type]['weights'] |
| bias_blob = lstm_pblobs[input_type]['biases'] |
| cur_weight = workspace.FetchBlob(weights_blob) |
| cur_biases = workspace.FetchBlob(bias_blob) |
| |
| workspace.FeedBlob( |
| weights_blob, |
| wmat.reshape(cur_weight.shape).astype(np.float32)) |
| workspace.FeedBlob( |
| bias_blob, |
| bm.reshape(cur_biases.shape).astype(np.float32)) |
| |
| |
| def cudnn_LSTM(model, input_blob, initial_states, dim_in, dim_out, |
| scope, recurrent_params=None, input_params=None, |
| num_layers=1, return_params=False): |
| ''' |
| CuDNN version of LSTM for GPUs. |
| input_blob Blob containing the input. Will need to be available |
| when param_init_net is run, because the sequence lengths |
| and batch sizes will be inferred from the size of this |
| blob. |
| initial_states tuple of (hidden_init, cell_init) blobs |
| dim_in input dimensions |
| dim_out output/hidden dimension |
| scope namescope to apply |
| recurrent_params dict of blobs containing values for recurrent |
| gate weights, biases (if None, use random init values) |
| See GetLSTMParamNames() for format. |
| input_params dict of blobs containing values for input |
| gate weights, biases (if None, use random init values) |
| See GetLSTMParamNames() for format. |
| num_layers number of LSTM layers |
| return_params if True, returns (param_extract_net, param_mapping) |
| where param_extract_net is a net that when run, will |
| populate the blobs specified in param_mapping with the |
| current gate weights and biases (input/recurrent). |
| Useful for assigning the values back to non-cuDNN |
| LSTM. |
| ''' |
| with core.NameScope(scope): |
| weight_params = GetLSTMParamNames()['weights'] |
| bias_params = GetLSTMParamNames()['biases'] |
| |
| input_weight_size = dim_out * dim_in |
| upper_layer_input_weight_size = dim_out * dim_out |
| recurrent_weight_size = dim_out * dim_out |
| input_bias_size = dim_out |
| recurrent_bias_size = dim_out |
| |
| def init(layer, pname, input_type): |
| input_weight_size_for_layer = input_weight_size if layer == 0 else \ |
| upper_layer_input_weight_size |
| if pname in weight_params: |
| sz = input_weight_size_for_layer if input_type == 'input' \ |
| else recurrent_weight_size |
| elif pname in bias_params: |
| sz = input_bias_size if input_type == 'input' \ |
| else recurrent_bias_size |
| else: |
| assert False, "unknown parameter type {}".format(pname) |
| return model.param_init_net.UniformFill( |
| [], |
| "lstm_init_{}_{}_{}".format(input_type, pname, layer), |
| shape=[sz]) |
| |
| # Multiply by 4 since we have 4 gates per LSTM unit |
| first_layer_sz = input_weight_size + recurrent_weight_size + \ |
| input_bias_size + recurrent_bias_size |
| upper_layer_sz = upper_layer_input_weight_size + \ |
| recurrent_weight_size + input_bias_size + \ |
| recurrent_bias_size |
| total_sz = 4 * (first_layer_sz + (num_layers - 1) * upper_layer_sz) |
| |
| weights = model.create_param( |
| 'lstm_weight', |
| shape=[total_sz], |
| initializer=Initializer('UniformFill'), |
| tags=ParameterTags.WEIGHT, |
| ) |
| |
| lstm_args = { |
| 'hidden_size': dim_out, |
| 'rnn_mode': 'lstm', |
| 'bidirectional': 0, # TODO |
| 'dropout': 1.0, # TODO |
| 'input_mode': 'linear', # TODO |
| 'num_layers': num_layers, |
| 'engine': 'CUDNN' |
| } |
| |
| param_extract_net = core.Net("lstm_param_extractor") |
| param_extract_net.AddExternalInputs([input_blob, weights]) |
| param_extract_mapping = {} |
| |
| # Populate the weights-blob from blobs containing parameters for |
| # the individual components of the LSTM, such as forget/input gate |
| # weights and bises. Also, create a special param_extract_net that |
| # can be used to grab those individual params from the black-box |
| # weights blob. These results can be then fed to InitFromLSTMParams() |
| for input_type in ['input', 'recurrent']: |
| param_extract_mapping[input_type] = {} |
| p = recurrent_params if input_type == 'recurrent' else input_params |
| if p is None: |
| p = {} |
| for pname in weight_params + bias_params: |
| for j in range(0, num_layers): |
| values = p[pname] if pname in p else init(j, pname, input_type) |
| model.param_init_net.RecurrentParamSet( |
| [input_blob, weights, values], |
| weights, |
| layer=j, |
| input_type=input_type, |
| param_type=pname, |
| **lstm_args |
| ) |
| if pname not in param_extract_mapping[input_type]: |
| param_extract_mapping[input_type][pname] = {} |
| b = param_extract_net.RecurrentParamGet( |
| [input_blob, weights], |
| ["lstm_{}_{}_{}".format(input_type, pname, j)], |
| layer=j, |
| input_type=input_type, |
| param_type=pname, |
| **lstm_args |
| ) |
| param_extract_mapping[input_type][pname][j] = b |
| |
| (hidden_input_blob, cell_input_blob) = initial_states |
| output, hidden_output, cell_output, rnn_scratch, dropout_states = \ |
| model.net.Recurrent( |
| [input_blob, hidden_input_blob, cell_input_blob, weights], |
| ["lstm_output", "lstm_hidden_output", "lstm_cell_output", |
| "lstm_rnn_scratch", "lstm_dropout_states"], |
| seed=random.randint(0, 100000), # TODO: dropout seed |
| **lstm_args |
| ) |
| model.net.AddExternalOutputs( |
| hidden_output, cell_output, rnn_scratch, dropout_states) |
| |
| if return_params: |
| param_extract = param_extract_net, param_extract_mapping |
| return output, hidden_output, cell_output, param_extract |
| else: |
| return output, hidden_output, cell_output |
| |
| |
| def LSTMWithAttention( |
| model, |
| decoder_inputs, |
| decoder_input_lengths, |
| initial_decoder_hidden_state, |
| initial_decoder_cell_state, |
| initial_attention_weighted_encoder_context, |
| encoder_output_dim, |
| encoder_outputs, |
| encoder_lengths, |
| decoder_input_dim, |
| decoder_state_dim, |
| scope, |
| attention_type=AttentionType.Regular, |
| outputs_with_grads=(0, 4), |
| weighted_encoder_outputs=None, |
| lstm_memory_optimization=False, |
| attention_memory_optimization=False, |
| forget_bias=0.0, |
| forward_only=False, |
| ): |
| ''' |
| Adds a LSTM with attention mechanism to a model. |
| |
| The implementation is based on https://arxiv.org/abs/1409.0473, with |
| a small difference in the order |
| how we compute new attention context and new hidden state, similarly to |
| https://arxiv.org/abs/1508.04025. |
| |
| The model uses encoder-decoder naming conventions, |
| where the decoder is the sequence the op is iterating over, |
| while computing the attention context over the encoder. |
| |
| model: ModelHelper object new operators would be added to |
| |
| decoder_inputs: the input sequence in a format T x N x D |
| where T is sequence size, N - batch size and D - input dimension |
| |
| decoder_input_lengths: blob containing sequence lengths |
| which would be passed to LSTMUnit operator |
| |
| initial_decoder_hidden_state: initial hidden state of LSTM |
| |
| initial_decoder_cell_state: initial cell state of LSTM |
| |
| initial_attention_weighted_encoder_context: initial attention context |
| |
| encoder_output_dim: dimension of encoder outputs |
| |
| encoder_outputs: the sequence, on which we compute the attention context |
| at every iteration |
| |
| encoder_lengths: a tensor with lengths of each encoder sequence in batch |
| (may be None, meaning all encoder sequences are of same length) |
| |
| decoder_input_dim: input dimension (last dimension on decoder_inputs) |
| |
| decoder_state_dim: size of hidden states of LSTM |
| |
| attention_type: One of: AttentionType.Regular, AttentionType.Recurrent. |
| Determines which type of attention mechanism to use. |
| |
| outputs_with_grads : position indices of output blobs which will receive |
| external error gradient during backpropagation |
| |
| weighted_encoder_outputs: encoder outputs to be used to compute attention |
| weights. In the basic case it's just linear transformation of |
| encoder outputs (that the default, when weighted_encoder_outputs is None). |
| However, it can be something more complicated - like a separate |
| encoder network (for example, in case of convolutional encoder) |
| |
| lstm_memory_optimization: recompute LSTM activations on backward pass, so |
| we don't need to store their values in forward passes |
| |
| attention_memory_optimization: recompute attention for backward pass |
| |
| forward_only: whether to create only forward pass |
| ''' |
| cell = LSTMWithAttentionCell( |
| encoder_output_dim=encoder_output_dim, |
| encoder_outputs=encoder_outputs, |
| encoder_lengths=encoder_lengths, |
| decoder_input_dim=decoder_input_dim, |
| decoder_state_dim=decoder_state_dim, |
| name=scope, |
| attention_type=attention_type, |
| weighted_encoder_outputs=weighted_encoder_outputs, |
| forget_bias=forget_bias, |
| lstm_memory_optimization=lstm_memory_optimization, |
| attention_memory_optimization=attention_memory_optimization, |
| forward_only=forward_only, |
| ) |
| initial_states = [ |
| initial_decoder_hidden_state, |
| initial_decoder_cell_state, |
| initial_attention_weighted_encoder_context, |
| ] |
| if attention_type == AttentionType.SoftCoverage: |
| initial_states.append(cell.build_initial_coverage(model)) |
| _, result = cell.apply_over_sequence( |
| model=model, |
| inputs=decoder_inputs, |
| seq_lengths=decoder_input_lengths, |
| initial_states=initial_states, |
| outputs_with_grads=outputs_with_grads, |
| ) |
| return result |
| |
| |
| def _layered_LSTM( |
| model, input_blob, seq_lengths, initial_states, |
| dim_in, dim_out, scope, outputs_with_grads=(0,), return_params=False, |
| memory_optimization=False, forget_bias=0.0, forward_only=False, |
| drop_states=False, create_lstm=None): |
| params = locals() # leave it as a first line to grab all params |
| params.pop('create_lstm') |
| if not isinstance(dim_out, list): |
| return create_lstm(**params) |
| elif len(dim_out) == 1: |
| params['dim_out'] = dim_out[0] |
| return create_lstm(**params) |
| |
| assert len(dim_out) != 0, "dim_out list can't be empty" |
| assert return_params is False, "return_params not supported for layering" |
| for i, output_dim in enumerate(dim_out): |
| params.update({ |
| 'dim_out': output_dim |
| }) |
| output, last_output, all_states, last_state = create_lstm(**params) |
| params.update({ |
| 'input_blob': output, |
| 'dim_in': output_dim, |
| 'initial_states': (last_output, last_state), |
| 'scope': scope + '_layer_{}'.format(i + 1) |
| }) |
| return output, last_output, all_states, last_state |
| |
| |
| layered_LSTM = functools.partial(_layered_LSTM, create_lstm=LSTM) |
