def call(self, inputs, mask=None, training=None, initial_state=None, constants=None): # The input should be dense, padded with zeros. If a ragged input is fed # into the layer, it is padded and the row lengths are used for masking. inputs, row_lengths = backend.convert_inputs_if_ragged(inputs) is_ragged_input = (row_lengths is not None) self._validate_args_if_ragged(is_ragged_input, mask) inputs, initial_state, constants = self._process_inputs( inputs, initial_state, constants) self._maybe_reset_cell_dropout_mask(self.cell) if isinstance(self.cell, StackedRNNCells): for cell in self.cell.cells: self._maybe_reset_cell_dropout_mask(cell) if mask is not None: # Time step masks must be the same for each input. # TODO(scottzhu): Should we accept multiple different masks? mask = tf.nest.flatten(mask)[0] if tf.nest.is_nested(inputs): # In the case of nested input, use the first element for shape check. input_shape = backend.int_shape(tf.nest.flatten(inputs)[0]) else: input_shape = backend.int_shape(inputs) timesteps = input_shape[0] if self.time_major else input_shape[1] if self.unroll and timesteps is None: raise ValueError('Cannot unroll a RNN if the ' 'time dimension is undefined. \n' '- If using a Sequential model, ' 'specify the time dimension by passing ' 'an `input_shape` or `batch_input_shape` ' 'argument to your first layer. If your ' 'first layer is an Embedding, you can ' 'also use the `input_length` argument.\n' '- If using the functional API, specify ' 'the time dimension by passing a `shape` ' 'or `batch_shape` argument to your Input layer.') kwargs = {} if generic_utils.has_arg(self.cell.call, 'training'): kwargs['training'] = training # TF RNN cells expect single tensor as state instead of list wrapped tensor. is_tf_rnn_cell = getattr(self.cell, '_is_tf_rnn_cell', None) is not None # Use the __call__ function for callable objects, eg layers, so that it # will have the proper name scopes for the ops, etc. cell_call_fn = self.cell.__call__ if callable( self.cell) else self.cell.call if constants: if not generic_utils.has_arg(self.cell.call, 'constants'): raise ValueError( f'RNN cell {self.cell} does not support constants. ' f'Received: constants={constants}') def step(inputs, states): constants = states[-self._num_constants:] # pylint: disable=invalid-unary-operand-type states = states[:-self._num_constants] # pylint: disable=invalid-unary-operand-type states = states[0] if len( states) == 1 and is_tf_rnn_cell else states output, new_states = cell_call_fn(inputs, states, constants=constants, **kwargs) if not tf.nest.is_nested(new_states): new_states = [new_states] return output, new_states else: def step(inputs, states): states = states[0] if len( states) == 1 and is_tf_rnn_cell else states output, new_states = cell_call_fn(inputs, states, **kwargs) if not tf.nest.is_nested(new_states): new_states = [new_states] return output, new_states last_output, outputs, states = backend.rnn( step, inputs, initial_state, constants=constants, go_backwards=self.go_backwards, mask=mask, unroll=self.unroll, input_length=row_lengths if row_lengths is not None else timesteps, time_major=self.time_major, zero_output_for_mask=self.zero_output_for_mask) if self.stateful: updates = [ tf.compat.v1.assign(self_state, tf.cast(state, self_state.dtype)) for self_state, state in zip(tf.nest.flatten(self.states), tf.nest.flatten(states)) ] self.add_update(updates) if self.return_sequences: output = backend.maybe_convert_to_ragged( is_ragged_input, outputs, row_lengths, go_backwards=self.go_backwards) else: output = last_output if self.return_state: if not isinstance(states, (list, tuple)): states = [states] else: states = list(states) return generic_utils.to_list(output) + states else: return output
def call(self, inputs, mask=None, training=None, initial_state=None): # The input should be dense, padded with zeros. If a ragged input is fed # into the layer, it is padded and the row lengths are used for masking. inputs, row_lengths = backend.convert_inputs_if_ragged(inputs) is_ragged_input = row_lengths is not None self._validate_args_if_ragged(is_ragged_input, mask) # GRU does not support constants. Ignore it during process. inputs, initial_state, _ = self._process_inputs( inputs, initial_state, None) if isinstance(mask, list): mask = mask[0] input_shape = backend.int_shape(inputs) timesteps = input_shape[0] if self.time_major else input_shape[1] if not self._could_use_gpu_kernel: kwargs = {"training": training} self._maybe_reset_cell_dropout_mask(self.cell) def step(cell_inputs, cell_states): return self.cell(cell_inputs, cell_states, **kwargs) last_output, outputs, states = backend.rnn( step, inputs, initial_state, constants=None, go_backwards=self.go_backwards, mask=mask, unroll=self.unroll, input_length=row_lengths if row_lengths is not None else timesteps, time_major=self.time_major, zero_output_for_mask=self.zero_output_for_mask, return_all_outputs=self.return_sequences, ) # This is a dummy tensor for testing purpose. runtime = gru_lstm_utils.runtime(gru_lstm_utils.RUNTIME_UNKNOWN) else: last_output, outputs, runtime, states = self._defun_gru_call( inputs, initial_state, training, mask, row_lengths) if self.stateful: updates = [ tf.compat.v1.assign(self.states[0], tf.cast(states[0], self.states[0].dtype)) ] self.add_update(updates) if self.return_sequences: output = backend.maybe_convert_to_ragged( is_ragged_input, outputs, row_lengths, go_backwards=self.go_backwards, ) else: output = last_output if self.return_state: return [output] + list(states) elif self._return_runtime: return output, runtime else: return output
def call(self, inputs, training=None, mask=None): kwargs = {} if generic_utils.has_arg(self.layer.call, 'training'): kwargs['training'] = training input_shape = tf.nest.map_structure( lambda x: tf.TensorShape(K.int_shape(x)), inputs) batch_size = tf_utils.convert_shapes(input_shape) batch_size = tf.nest.flatten(batch_size)[0] if batch_size and not self._always_use_reshape: inputs, row_lengths = K.convert_inputs_if_ragged(inputs) is_ragged_input = row_lengths is not None input_length = tf_utils.convert_shapes(input_shape) input_length = tf.nest.flatten(input_length)[1] # batch size matters, use rnn-based implementation def step(x, _): output = self.layer(x, **kwargs) return output, [] _, outputs, _ = K.rnn( step, inputs, initial_states=[], input_length=row_lengths[0] if is_ragged_input else input_length, mask=mask, unroll=False) # pylint: disable=g-long-lambda y = tf.nest.map_structure( lambda output: K.maybe_convert_to_ragged(is_ragged_input, output, row_lengths), outputs) else: # No batch size specified, therefore the layer will be able # to process batches of any size. # We can go with reshape-based implementation for performance. is_ragged_input = tf.nest.map_structure( lambda x: isinstance(x, tf.RaggedTensor), inputs) is_ragged_input = tf.nest.flatten(is_ragged_input) if all(is_ragged_input): input_values = tf.nest.map_structure(lambda x: x.values, inputs) input_row_lenghts = tf.nest.map_structure( lambda x: x.nested_row_lengths()[0], inputs) y = self.layer(input_values, **kwargs) y = tf.nest.map_structure(tf.RaggedTensor.from_row_lengths, y, input_row_lenghts) elif any(is_ragged_input): raise ValueError('All inputs has to be either ragged or not, ' 'but not mixed. You passed: {}'.format(inputs)) else: input_length = tf_utils.convert_shapes(input_shape) input_length = tf.nest.flatten(input_length)[1] if not input_length: input_length = tf.nest.map_structure(lambda x: tf.compat.v1.shape(x)[1], inputs) input_length = generic_utils.to_list(tf.nest.flatten(input_length))[0] inner_input_shape = tf.nest.map_structure( lambda x: self._get_shape_tuple((-1,), x, 2), inputs) # Shape: (num_samples * timesteps, ...). And track the # transformation in self._input_map. inputs = tf.__internal__.nest.map_structure_up_to(inputs, tf.reshape, inputs, inner_input_shape) # (num_samples * timesteps, ...) if generic_utils.has_arg(self.layer.call, 'mask') and mask is not None: inner_mask_shape = self._get_shape_tuple((-1,), mask, 2) kwargs['mask'] = K.reshape(mask, inner_mask_shape) y = self.layer(inputs, **kwargs) # Shape: (num_samples, timesteps, ...) output_shape = self.compute_output_shape(input_shape) # pylint: disable=g-long-lambda output_shape = tf.nest.map_structure( lambda tensor, int_shape: self._get_shape_tuple( (-1, input_length), tensor, 1, int_shape[2:]), y, output_shape) y = tf.__internal__.nest.map_structure_up_to(y, tf.reshape, y, output_shape) if not tf.executing_eagerly(): # Set the static shape for the result since it might be lost during # array_ops reshape, eg, some `None` dim in the result could be # inferred. tf.__internal__.nest.map_structure_up_to( y, lambda tensor, shape: tensor.set_shape(shape), y, self.compute_output_shape(input_shape)) return y
def call(self, inputs, mask=None, training=None, initial_state=None): # The input should be dense, padded with zeros. If a ragged input is fed # into the layer, it is padded and the row lengths are used for masking. inputs, row_lengths = backend.convert_inputs_if_ragged(inputs) is_ragged_input = (row_lengths is not None) self._validate_args_if_ragged(is_ragged_input, mask) # LSTM does not support constants. Ignore it during process. inputs, initial_state, _ = self._process_inputs(inputs, initial_state, None) if isinstance(mask, list): mask = mask[0] input_shape = backend.int_shape(inputs) timesteps = input_shape[0] if self.time_major else input_shape[1] if not self._could_use_gpu_kernel: # Fall back to use the normal LSTM. kwargs = {'training': training} self._maybe_reset_cell_dropout_mask(self.cell) def step(inputs, states): return self.cell(inputs, states, **kwargs) last_output, outputs, states = backend.rnn( step, inputs, initial_state, constants=None, go_backwards=self.go_backwards, mask=mask, unroll=self.unroll, input_length=row_lengths if row_lengths is not None else timesteps, time_major=self.time_major, zero_output_for_mask=self.zero_output_for_mask) runtime = gru_lstm_utils.runtime(gru_lstm_utils.RUNTIME_UNKNOWN) else: # Use the new defun approach for backend implementation swap. # Note that different implementations need to have same function # signature, eg, the tensor parameters need to have same shape and dtypes. # Since the cuDNN has an extra set of bias, those bias will be passed to # both normal and cuDNN implementations. self.reset_dropout_mask() dropout_mask = self.get_dropout_mask_for_cell(inputs, training, count=4) if dropout_mask is not None: inputs = inputs * dropout_mask[0] if gru_lstm_utils.use_new_gru_lstm_impl(): lstm_kwargs = { 'inputs': inputs, 'init_h': gru_lstm_utils.read_variable_value(initial_state[0]), 'init_c': gru_lstm_utils.read_variable_value(initial_state[1]), 'kernel': gru_lstm_utils.read_variable_value(self.cell.kernel), 'recurrent_kernel': gru_lstm_utils.read_variable_value(self.cell.recurrent_kernel), 'bias': gru_lstm_utils.read_variable_value(self.cell.bias), 'mask': mask, 'time_major': self.time_major, 'go_backwards': self.go_backwards, 'sequence_lengths': row_lengths, 'zero_output_for_mask': self.zero_output_for_mask, } (last_output, outputs, new_h, new_c, runtime) = self._defun_wrapper.defun_layer(**lstm_kwargs) else: gpu_lstm_kwargs = { 'inputs': inputs, 'init_h': gru_lstm_utils.read_variable_value(initial_state[0]), 'init_c': gru_lstm_utils.read_variable_value(initial_state[1]), 'kernel': gru_lstm_utils.read_variable_value(self.cell.kernel), 'recurrent_kernel': gru_lstm_utils.read_variable_value(self.cell.recurrent_kernel), 'bias': gru_lstm_utils.read_variable_value(self.cell.bias), 'mask': mask, 'time_major': self.time_major, 'go_backwards': self.go_backwards, 'sequence_lengths': row_lengths } normal_lstm_kwargs = gpu_lstm_kwargs.copy() normal_lstm_kwargs.update({ 'zero_output_for_mask': self.zero_output_for_mask, }) if tf.executing_eagerly(): device_type = gru_lstm_utils.get_context_device_type() can_use_gpu = ( # Either user specified GPU or unspecified but GPU is available. (device_type == gru_lstm_utils.GPU_DEVICE_NAME or (device_type is None and tf.config.list_logical_devices('GPU'))) and (mask is None or gru_lstm_utils.is_cudnn_supported_inputs(mask, self.time_major))) # Under eager context, check the device placement and prefer the # GPU implementation when GPU is available. if can_use_gpu: last_output, outputs, new_h, new_c, runtime = gpu_lstm( **gpu_lstm_kwargs) else: last_output, outputs, new_h, new_c, runtime = standard_lstm( **normal_lstm_kwargs) else: (last_output, outputs, new_h, new_c, runtime) = lstm_with_backend_selection(**normal_lstm_kwargs) states = [new_h, new_c] if self.stateful: updates = [ tf.compat.v1.assign(self_state, tf.cast(state, self_state.dtype)) for self_state, state in zip(self.states, states) ] self.add_update(updates) if self.return_sequences: output = backend.maybe_convert_to_ragged( is_ragged_input, outputs, row_lengths, go_backwards=self.go_backwards) else: output = last_output if self.return_state: return [output] + list(states) elif self.return_runtime: return output, runtime else: return output