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()
