def __init__(self, memory_size=128, word_size=20, num_reads=1, num_writes=1, name='memory_access'): """Creates a MemoryAccess module. Args: memory_size: The number of memory slots (N in the DNC paper). word_size: The width of each memory slot (W in the DNC paper) num_reads: The number of read heads (R in the DNC paper). num_writes: The number of write heads (fixed at 1 in the paper). name: The name of the module. """ super(MemoryAccess, self).__init__(name=name) self._memory_size = memory_size self._word_size = word_size self._num_reads = num_reads self._num_writes = num_writes self._write_content_weights_mod = addressing.CosineWeights( num_writes, word_size, name='write_content_weights') self._read_content_weights_mod = addressing.CosineWeights( num_reads, word_size, name='read_content_weights') self._linkage = addressing.TemporalLinkage(memory_size, num_writes) self._freeness = addressing.Freeness(memory_size)
def testWriteAllocationWeights(self): batch_size = 7 memory_size = 23 num_writes = 5 module = addressing.Freeness(memory_size) usage = np.random.rand(batch_size, memory_size) write_gates = np.random.rand(batch_size, num_writes) # Turn off gates for heads 1 and 3 in batch 0. This doesn't scaling down the # weighting, but it means that the usage doesn't change, so we should get # the same allocation weightings for: (1, 2) and (3, 4) (but all others # being different). write_gates[0, 1] = 0 write_gates[0, 3] = 0 # and turn heads 0 and 2 on for full effect. write_gates[0, 0] = 1 write_gates[0, 2] = 1 # In batch 1, make one of the usages 0 and another almost 0, so that these # entries get most of the allocation weights for the first and second heads. usage[ 1] = usage[1] * 0.9 + 0.1 # make sure all entries are in [0.1, 1] usage[1][4] = 0 # write head 0 should get allocated to position 4 usage[1][3] = 1e-4 # write head 1 should get allocated to position 3 write_gates[1, 0] = 1 # write head 0 fully on write_gates[1, 1] = 1 # write head 1 fully on weights = module.write_allocation_weights( usage=tf.constant(usage), write_gates=tf.constant(write_gates), num_writes=num_writes) with self.test_session(): weights = weights.eval() # Check that all weights are between 0 and 1 self.assertGreaterEqual(weights.min(), 0) self.assertLessEqual(weights.max(), 1) # Check that weights sum to close to 1 self.assertAllClose(np.sum(weights, axis=2), np.ones([batch_size, num_writes]), atol=1e-3) # Check the same / different allocation weight pairs as described above. self.assertGreater( np.abs(weights[0, 0, :] - weights[0, 1, :]).max(), 0.1) self.assertAllEqual(weights[0, 1, :], weights[0, 2, :]) self.assertGreater( np.abs(weights[0, 2, :] - weights[0, 3, :]).max(), 0.1) self.assertAllEqual(weights[0, 3, :], weights[0, 4, :]) self.assertAllClose(weights[1][0], util.one_hot(memory_size, 4), atol=1e-3) self.assertAllClose(weights[1][1], util.one_hot(memory_size, 3), atol=1e-3)
def testAllocationGradient(self): batch_size = 1 memory_size = 5 usage = tf.constant(np.random.rand(batch_size, memory_size)) module = addressing.Freeness(memory_size) allocation = module._allocation(usage) with self.test_session(): err = tf.test.compute_gradient_error( usage, usage.get_shape().as_list(), allocation, allocation.get_shape().as_list(), delta=1e-5) self.assertLess(err, 0.01)
def testAllocation(self): batch_size = 7 memory_size = 13 usage = np.random.rand(batch_size, memory_size) module = addressing.Freeness(memory_size) allocation = module._allocation(tf.constant(usage)) with self.test_session(): allocation = allocation.eval() # 1. Test that max allocation goes to min usage, and vice versa. self.assertAllEqual(np.argmin(usage, axis=1), np.argmax(allocation, axis=1)) self.assertAllEqual(np.argmax(usage, axis=1), np.argmin(allocation, axis=1)) # 2. Test that allocations sum to almost 1. self.assertAllClose(np.sum(allocation, axis=1), np.ones(batch_size), 0.01)
def testWriteAllocationWeightsGradient(self): batch_size = 7 memory_size = 5 num_writes = 3 module = addressing.Freeness(memory_size) usage = tf.constant(np.random.rand(batch_size, memory_size)) write_gates = tf.constant(np.random.rand(batch_size, num_writes)) weights = module.write_allocation_weights(usage, write_gates, num_writes) with self.test_session(): err = tf.test.compute_gradient_error( [usage, write_gates], [usage.get_shape().as_list(), write_gates.get_shape().as_list()], weights, weights.get_shape().as_list(), delta=1e-5) self.assertLess(err, 0.01)
def testModule(self): batch_size = 5 memory_size = 11 num_reads = 3 num_writes = 7 module = addressing.Freeness(memory_size) free_gate = np.random.rand(batch_size, num_reads) # Produce read weights that sum to 1 for each batch and head. prev_read_weights = np.random.rand(batch_size, num_reads, memory_size) prev_read_weights[1, :, 3] = 0 # no read at batch 1, position 3; see below prev_read_weights /= prev_read_weights.sum(2, keepdims=True) prev_write_weights = np.random.rand(batch_size, num_writes, memory_size) prev_write_weights /= prev_write_weights.sum(2, keepdims=True) prev_usage = np.random.rand(batch_size, memory_size) # Add some special values that allows us to test the behaviour: prev_write_weights[1, 2, 3] = 1 # full write in batch 1, head 2, position 3 prev_read_weights[2, 0, 4] = 1 # full read at batch 2, head 0, position 4 free_gate[2, 0] = 1 # can free up all locations for batch 2, read head 0 usage = module(tf.constant(prev_write_weights), tf.constant(free_gate), tf.constant(prev_read_weights), tf.constant(prev_usage)) with self.test_session(): usage = usage.eval() # Check all usages are between 0 and 1. self.assertGreaterEqual(usage.min(), 0) self.assertLessEqual(usage.max(), 1) # Check that the full write at batch 1, position 3 makes it fully used. self.assertEqual(usage[1][3], 1) # Check that the full free at batch 2, position 4 makes it fully free. self.assertEqual(usage[2][4], 0)