# Lint as: python2, python3 # pylint: disable=g-bad-file-header # Copyright 2019 DeepMind Technologies Limited. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ """Memory Reader/Writer for RMA.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import collections import sonnet as snt import tensorflow.compat.v1 as tf ReadInformation = collections.namedtuple( 'ReadInformation', ('weights', 'indices', 'keys', 'strengths')) class MemoryWriter(snt.RNNCore): """Memory Writer Module.""" def __init__(self, mem_shape, name='memory_writer'): """Initializes the `MemoryWriter`. Args: mem_shape: The shape of the memory `(num_rows, memory_width)`. name: The name to use for the Sonnet module. """ super(MemoryWriter, self).__init__(name=name) self._mem_shape = mem_shape def _build(self, inputs, state): """Inserts z into the argmin row of usage markers and updates all rows. Returns an operation that, when executed, correctly updates the internal state and usage markers. Args: inputs: A tuple consisting of: * z, the value to write at this timestep * mem_state, the state of the memory at this timestep before writing state: The state is just the write_counter. Returns: A tuple of the new memory state and a tuple containing the next state. """ z, mem_state = inputs # Stop gradient on writes to memory. z = tf.stop_gradient(z) prev_write_counter = state new_row_value = z # Find the index to insert the next row into. num_mem_rows = self._mem_shape[0] write_index = tf.cast(prev_write_counter, dtype=tf.int32) % num_mem_rows one_hot_row = tf.one_hot(write_index, num_mem_rows) write_counter = prev_write_counter + 1 # Insert state variable to new row. # First you need to size it up to the full size. insert_new_row = lambda mem, o_hot, z: mem - (o_hot * mem) + (o_hot * z) new_mem = insert_new_row(mem_state, tf.expand_dims(one_hot_row, axis=-1), tf.expand_dims(new_row_value, axis=-2)) new_state = write_counter return new_mem, new_state @property def state_size(self): """Returns a description of the state size, without batch dimension.""" return tf.TensorShape([]) @property def output_size(self): """Returns a description of the output size, without batch dimension.""" return self._mem_shape class MemoryReader(snt.AbstractModule): """Memory Reader Module.""" def __init__(self, memory_word_size, num_read_heads, top_k=0, memory_size=None, name='memory_reader'): """Initializes the `MemoryReader`. Args: memory_word_size: The dimension of the 1-D read keys this memory reader should produce. Each row of the memory is of length `memory_word_size`. num_read_heads: The number of reads to perform. top_k: Softmax and summation when reading is only over top k most similar entries in memory. top_k=0 (default) means dense reads, i.e. no top_k. memory_size: Number of rows in memory. name: The name for this Sonnet module. """ super(MemoryReader, self).__init__(name=name) self._memory_word_size = memory_word_size self._num_read_heads = num_read_heads self._top_k = top_k # This is not an RNNCore but it is useful to expose the output size. self._output_size = num_read_heads * memory_word_size num_read_weights = top_k if top_k > 0 else memory_size self._read_info_size = ReadInformation( weights=tf.TensorShape([num_read_heads, num_read_weights]), indices=tf.TensorShape([num_read_heads, num_read_weights]), keys=tf.TensorShape([num_read_heads, memory_word_size]), strengths=tf.TensorShape([num_read_heads]), ) with self._enter_variable_scope(): # Transforms to value-based read for each read head. output_dim = (memory_word_size + 1) * num_read_heads self._keys_and_read_strengths_generator = snt.Linear(output_dim) def _build(self, inputs): """Looks up rows in memory. In the args list, we have the following conventions: B: batch size M: number of slots in a row of the memory matrix R: number of rows in the memory matrix H: number of read heads in the memory controller Args: inputs: A tuple of * read_inputs, a tensor of shape [B, ...] that will be flattened and passed through a linear layer to get read keys/read_strengths for each head. * mem_state, the primary memory tensor. Of shape [B, R, M]. Returns: The read from the memory (concatenated across read heads) and read information. """ # Assert input shapes are compatible and separate inputs. _assert_compatible_memory_reader_input(inputs) read_inputs, mem_state = inputs # Determine the read weightings for each key. flat_outputs = self._keys_and_read_strengths_generator( snt.BatchFlatten()(read_inputs)) # Separate the read_strengths from the rest of the weightings. h = self._num_read_heads flat_keys = flat_outputs[:, :-h] read_strengths = tf.nn.softplus(flat_outputs[:, -h:]) # Reshape the weights. read_shape = (self._num_read_heads, self._memory_word_size) read_keys = snt.BatchReshape(read_shape)(flat_keys) # Read from memory. memory_reads, read_weights, read_indices, read_strengths = ( read_from_memory(read_keys, read_strengths, mem_state, self._top_k)) concatenated_reads = snt.BatchFlatten()(memory_reads) return concatenated_reads, ReadInformation( weights=read_weights, indices=read_indices, keys=read_keys, strengths=read_strengths) @property def output_size(self): """Returns a description of the output size, without batch dimension.""" return self._output_size, self._read_info_size def read_from_memory(read_keys, read_strengths, mem_state, top_k): """Function for cosine similarity content based reading from memory matrix. In the args list, we have the following conventions: B: batch size M: number of slots in a row of the memory matrix R: number of rows in the memory matrix H: number of read heads (of the controller or the policy) K: top_k if top_k>0 Args: read_keys: the read keys of shape [B, H, M]. read_strengths: the coefficients used to compute the normalised weighting vector of shape [B, H]. mem_state: the primary memory tensor. Of shape [B, R, M]. top_k: only use top k read matches, other reads do not go into softmax and are zeroed out in the output. top_k=0 (default) means use dense reads. Returns: The memory reads [B, H, M], read weights [B, H, top k], read indices [B, H, top k], and read strengths [B, H, 1]. """ _assert_compatible_read_from_memory_inputs(read_keys, read_strengths, mem_state) batch_size = read_keys.shape[0] num_read_heads = read_keys.shape[1] with tf.name_scope('memory_reading'): # Scale such that all rows are L2-unit vectors, for memory and read query. scaled_read_keys = tf.math.l2_normalize(read_keys, axis=-1) # [B, H, M] scaled_mem = tf.math.l2_normalize(mem_state, axis=-1) # [B, R, M] # The cosine distance is then their dot product. # Find the cosine distance between each read head and each row of memory. cosine_distances = tf.matmul( scaled_read_keys, scaled_mem, transpose_b=True) # [B, H, R] # The rank must match cosine_distances for broadcasting to work. read_strengths = tf.expand_dims(read_strengths, axis=-1) # [B, H, 1] weighted_distances = read_strengths * cosine_distances # [B, H, R] if top_k: # Get top k indices (row indices with top k largest weighted distances). top_k_output = tf.nn.top_k(weighted_distances, top_k, sorted=False) read_indices = top_k_output.indices # [B, H, K] # Create a sub-memory for each read head with only the top k rows. # Each batch_gather is [B, K, M] and the list stacks to [B, H, K, M]. topk_mem_per_head = [tf.batch_gather(mem_state, ri_this_head) for ri_this_head in tf.unstack(read_indices, axis=1)] topk_mem = tf.stack(topk_mem_per_head, axis=1) # [B, H, K, M] topk_scaled_mem = tf.math.l2_normalize(topk_mem, axis=-1) # [B, H, K, M] # Calculate read weights for each head's top k sub-memory. expanded_scaled_read_keys = tf.expand_dims( scaled_read_keys, axis=2) # [B, H, 1, M] topk_cosine_distances = tf.reduce_sum( expanded_scaled_read_keys * topk_scaled_mem, axis=-1) # [B, H, K] topk_weighted_distances = ( read_strengths * topk_cosine_distances) # [B, H, K] read_weights = tf.nn.softmax( topk_weighted_distances, axis=-1) # [B, H, K] # For each head, read using the sub-memories and corresponding weights. expanded_weights = tf.expand_dims(read_weights, axis=-1) # [B, H, K, 1] memory_reads = tf.reduce_sum( expanded_weights * topk_mem, axis=2) # [B, H, M] else: read_weights = tf.nn.softmax(weighted_distances, axis=-1) num_rows_memory = mem_state.shape[1] all_indices = tf.range(num_rows_memory, dtype=tf.int32) all_indices = tf.reshape(all_indices, [1, 1, num_rows_memory]) read_indices = tf.tile(all_indices, [batch_size, num_read_heads, 1]) # This is the actual memory access. # Note that matmul automatically batch applies for us. memory_reads = tf.matmul(read_weights, mem_state) read_keys.shape.assert_is_compatible_with(memory_reads.shape) read_strengths = tf.squeeze(read_strengths, axis=-1) # [B, H, 1] -> [B, H] return memory_reads, read_weights, read_indices, read_strengths def _assert_compatible_read_from_memory_inputs(read_keys, read_strengths, mem_state): read_keys.shape.assert_has_rank(3) b_shape, h_shape, m_shape = read_keys.shape mem_state.shape.assert_has_rank(3) r_shape = mem_state.shape[1] read_strengths.shape.assert_is_compatible_with( tf.TensorShape([b_shape, h_shape])) mem_state.shape.assert_is_compatible_with( tf.TensorShape([b_shape, r_shape, m_shape])) def _assert_compatible_memory_reader_input(input_tensors): """Asserts MemoryReader's _build has been given the correct shapes.""" assert len(input_tensors) == 2 _, mem_state = input_tensors mem_state.shape.assert_has_rank(3)