def test_eager_dnc_optimization(self): batch_size = 7 input_size = 15 memory_config = { 'memory_size': 27, 'word_size': 9, 'num_read_heads': 10, } output_size = 36 x = tf.keras.Input(shape=( None, input_size, )) dnc_cell = DNC(output_size, controller_units=30, **memory_config) dnc_initial_state = dnc_cell.get_initial_state(batch_size=batch_size) layer = tf.keras.layers.RNN(dnc_cell) y = layer(x, initial_state=dnc_initial_state) model = tf.keras.models.Model(x, y) model.compile(optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.001), loss='mse', run_eagerly=True) model.train_on_batch(np.zeros((batch_size, 5, input_size)), np.zeros((batch_size, output_size))) self.assertEqual(model.output_shape[1], output_size)
def test_constructor(self): memory_config = { 'memory_size': 4, 'word_size': 5, 'num_read_heads': 2, } dnc = DNC(10, controller_units=64, **memory_config) input_size = 17 test_input = np.random.uniform( -3, 3, (2, dnc._W * dnc._R + input_size)).astype(np.float32) initial_state = dnc.get_initial_state(batch_size=2) _, _ = dnc(test_input, initial_state) self.assertEqual(dnc._interface_vector_size, 38) self.assertEqual(dnc.output_size, 10) self.assertEqual(dnc.get_config()["name"], "DNC")
def test_dnc_output_shape(self): batch_size = 3 memory_config = { 'memory_size': 16, 'word_size': 5, 'num_read_heads': 7, } output_size = 10 for input_size in [10, 17, 49]: dnc = DNC(output_size, controller_units=64, **memory_config) initial_state = dnc.get_initial_state(batch_size=batch_size) input_shape = dnc._W * dnc._R + input_size test_input = np.random.uniform( -3, 3, (batch_size, input_shape)).astype(np.float32) example_output, _ = dnc( tf.convert_to_tensor(test_input), initial_state, ) self.assertEqual(example_output.shape, (batch_size, output_size))
import tensorflow as tf from dnc.dnc import DNC import numpy as np np.random.seed(1) g = tf.Graph() with g.as_default(): batch_size = 4 output_size = 20 input_size = 10 dnc = DNC(output_size, controller_units=128, memory_size=256, word_size=64, num_read_heads=4) initial_state = dnc.get_initial_state(batch_size=batch_size) example_input = np.random.uniform(0, 1, (batch_size, input_size)).astype( np.float32) output_op, _ = dnc( tf.convert_to_tensor(example_input), initial_state, ) init = tf.global_variables_initializer() with tf.Session(graph=g) as sess: init.run() example_output = sess.run(output_op) tf.summary.FileWriter("graphs", g).close()