def test_pre_build(Simulator):
class TestFunc:
def pre_build(self, size_in, size_out):
self.weights = tf.Variable(tf.ones((size_in[1], size_out[1])))
def __call__(self, t, x):
return tf.matmul(x, tf.cast(self.weights, x.dtype))
class TestFunc2:
def pre_build(self, size_in, size_out):
assert size_in is None
assert size_out == (1, 1)
def __call__(self, t):
return tf.reshape(t, (1, 1))
with nengo.Network() as net:
inp = nengo.Node([1, 1])
test = TensorNode(TestFunc(), size_in=2, size_out=3)
nengo.Connection(inp, test, synapse=None)
p = nengo.Probe(test)
test2 = TensorNode(TestFunc2(), size_out=1)
p2 = nengo.Probe(test2)
with Simulator(net) as sim:
sim.step()
assert np.allclose(sim.data[p], [[2, 2, 2]])
assert np.allclose(sim.data[p2], [[0.001]])
sim.reset()
sim.step()
assert np.allclose(sim.data[p], [[2, 2, 2]])
assert np.allclose(sim.data[p2], [[0.001]])
def test_pre_build(Simulator):
class TestLayer(tf.keras.layers.Layer):
def __init__(self, size_out):
super().__init__()
self.size_out = size_out
def build(self, shape_in):
super().build(shape_in)
self.w = self.add_weight(
initializer=tf.initializers.ones(),
shape=(shape_in[-1], self.size_out),
name="weights",
)
def call(self, x):
return tf.matmul(x, tf.cast(self.w, x.dtype))
def compute_output_shape(self, _):
return tf.TensorShape((None, self.size_out))
class TestLayer2(tf.keras.layers.Layer):
def build(self, shape_in):
# TODO: add this check back in once
# https://github.com/tensorflow/tensorflow/issues/32786 is fixed
# assert shape_in == [(), (1, 1, 1)]
pass
def call(self, inputs):
t, x = inputs
assert t.shape == ()
assert x.shape == (1, 1, 1)
return tf.reshape(t, (1, 1))
with nengo.Network() as net:
inp = nengo.Node([1, 1])
test = TensorNode(TestLayer(3), shape_in=(2, ), pass_time=False)
nengo.Connection(inp, test, synapse=None)
p = nengo.Probe(test)
test2 = TensorNode(TestLayer2(), shape_in=(1, 1), pass_time=True)
p2 = nengo.Probe(test2)
with Simulator(net) as sim:
sim.step()
assert np.allclose(sim.data[p], [[2, 2, 2]])
assert np.allclose(sim.data[p2], sim.trange()[:, None])
sim.reset()
sim.step()
assert np.allclose(sim.data[p], [[2, 2, 2]])
assert np.allclose(sim.data[p2], sim.trange()[:, None])
def test_validation():
with nengo.Network() as net:
with pytest.raises(ValidationError):
TensorNode([0])
with pytest.raises(ValidationError):
TensorNode(lambda t: t, size_out=0)
with pytest.raises(ValidationError):
TensorNode(lambda a, b, c: a)
with pytest.raises(ValidationError):
TensorNode(lambda x: None)
with pytest.raises(ValidationError):
TensorNode(lambda x: [0])
with pytest.raises(ValidationError):
TensorNode(lambda t: tf.zeros((2, 2, 2)))
n = TensorNode(lambda t: tf.zeros((5, 2)))
assert n.size_out == 2
n = TensorNode(lambda t: tf.zeros((5, 2)), size_out=4)
assert n.size_out == 4
with pytest.raises(BuildError):
with nengo.Simulator(net):
pass
def test_unspecified_shape():
with nengo.Network():
inp = nengo.Node([0])
with pytest.raises(ValidationError, match="must be an int"):
TensorNode(lambda x: x, shape_in=(None, 1))
with pytest.raises(ValidationError, match="must be an int"):
TensorNode(lambda x: x, shape_out=(None, 1))
with pytest.raises(ValidationError, match="must be an int"):
Layer(lambda x: x)(inp, shape_in=(None, 1))
with pytest.raises(ValidationError, match="must be an int"):
Layer(lambda x: x)(inp, shape_out=(None, 1))
def test_node(Simulator):
minibatch_size = 3
with nengo.Network() as net:
node0 = TensorNode(lambda t: tf.tile(tf.reshape(t, (1, -1)),
(minibatch_size, 1)))
node1 = TensorNode(lambda t, x: tf.sin(x), shape_in=(1, ))
nengo.Connection(node0, node1, synapse=None)
p0 = nengo.Probe(node0)
p1 = nengo.Probe(node1)
with Simulator(net, minibatch_size=minibatch_size) as sim:
sim.run_steps(10)
assert np.allclose(sim.data[p0], sim.trange()[None, :, None])
assert np.allclose(sim.data[p1], np.sin(sim.trange()[None, :, None]))
def test_training_arg(Simulator):
class TrainingLayer(tf.keras.layers.Layer):
def __init__(self, expected):
super().__init__()
self.expected = expected
def call(self, inputs, training=None):
tf.assert_equal(training, self.expected)
return tf.reshape(inputs, (1, 1))
with nengo.Network() as net:
node = TensorNode(TrainingLayer(expected=False),
shape_in=None,
shape_out=(1, ))
nengo.Probe(node)
with Simulator(net) as sim:
sim.predict(n_steps=10)
with Simulator(net) as sim:
sim.compile(optimizer=tf.optimizers.SGD(0), loss=tf.losses.mse)
node.tensor_func.expected = True
sim.fit(n_steps=10, y=np.zeros((1, 1, 1)))
with net:
configure_settings(learning_phase=True)
with Simulator(net) as sim:
sim.predict(n_steps=10)
def test_post_build(Simulator):
class TestFunc:
def pre_build(self, size_in, size_out):
self.weights = tf.Variable(tf.zeros((size_in[1], size_out[1])))
def post_build(self, sess, rng):
assert isinstance(rng, np.random.RandomState)
init_op = tf.assign(self.weights, tf.ones((2, 3)))
sess.run(init_op)
def __call__(self, t, x):
return tf.matmul(x, tf.cast(self.weights, x.dtype))
with nengo.Network() as net:
inp = nengo.Node([1, 1])
test = TensorNode(TestFunc(), size_in=2, size_out=3)
nengo.Connection(inp, test, synapse=None)
p = nengo.Probe(test)
with Simulator(net) as sim:
sim.step()
assert np.allclose(sim.data[p], [[2, 2, 2]])
sim.reset()
sim.step()
assert np.allclose(sim.data[p], [[2, 2, 2]])
def test_reuse_vars(Simulator):
def my_func(_, x):
# note: the control dependencies thing is due to some weird tensorflow
# issue with creating variables inside while loops
with tf.control_dependencies(None):
w = tf.get_variable("weights", initializer=tf.constant(2.0))
return x * tf.cast(w, x.dtype)
with nengo.Network() as net:
configure_settings(trainable=False)
inp = nengo.Node([1])
node = TensorNode(my_func, size_in=1)
p = nengo.Probe(node)
nengo.Connection(inp, node, synapse=None)
with Simulator(net, unroll_simulation=5) as sim:
sim.run_steps(5)
assert np.allclose(sim.data[p], 2)
with sim.tensor_graph.graph.as_default():
vars = tf.trainable_variables()
assert len(vars) == 1
assert vars[0].get_shape() == ()
assert sim.sess.run(vars[0]) == 2
def test_reuse_vars(Simulator, pytestconfig):
class MyLayer(tf.keras.layers.Layer):
def build(self, input_shape):
self.w = self.add_weight(initializer=tf.initializers.constant(2.0),
name="weights")
def call(self, x):
return x * tf.cast(self.w, x.dtype)
with nengo.Network() as net:
configure_settings(trainable=False)
inp = nengo.Node([1])
node = TensorNode(MyLayer(), shape_in=(1, ), pass_time=False)
nengo.Connection(inp, node, synapse=None)
node2 = Layer(
tf.keras.layers.Dense(
units=10,
use_bias=False,
kernel_initializer=tf.initializers.constant(3),
dtype=pytestconfig.getoption("--dtype"),
))(inp)
p = nengo.Probe(node)
p2 = nengo.Probe(node2)
with Simulator(net, unroll_simulation=5) as sim:
sim.run_steps(5)
assert np.allclose(sim.data[p], 2)
assert np.allclose(sim.data[p2], 3)
# note: when inference-only=True the weights will be marked as non-trainable
if sim.tensor_graph.inference_only:
assert len(sim.tensor_graph.saved_state) == 2
assert len(sim.keras_model.non_trainable_variables) == 2
assert len(sim.keras_model.trainable_variables) == 0
vars = sim.keras_model.non_trainable_variables
else:
assert len(sim.tensor_graph.saved_state) == 2
assert len(sim.keras_model.non_trainable_variables) == 0
assert len(sim.keras_model.trainable_variables) == 2
vars = sim.keras_model.trainable_variables
assert len(vars) == 2
assert vars[0].shape == ()
assert tf.keras.backend.get_value(vars[0]) == 2
assert vars[1].shape == (1, 10)
assert np.allclose(tf.keras.backend.get_value(vars[1]), 3)
def test_wrapped_model(Simulator):
inp0 = tf.keras.Input((1, ))
out0 = tf.keras.layers.Dense(units=10)(inp0)
model0 = tf.keras.Model(inp0, out0)
model0.compile(loss="mse", metrics=["accuracy"])
class KerasWrapper(tf.keras.layers.Layer):
def __init__(self, model):
super().__init__()
self.model = model
def build(self, input_shape):
super().build(input_shape)
self.model = tf.keras.models.clone_model(self.model)
# TODO: this shouldn't be necessary once we're running natively in eager
# mode
if not hasattr(self.model, "_metrics_lock"):
self.model._metrics_lock = threading.Lock()
for layer in self.model.layers:
layer._metrics_lock = threading.Lock()
def call(self, inputs):
return self.model(inputs)
with nengo.Network() as net:
layer = KerasWrapper(model0)
keras_node = TensorNode(layer,
shape_in=(1, ),
shape_out=(10, ),
pass_time=False)
nengo.Probe(keras_node)
with Simulator(net) as sim:
# this caused an error at one point, so testing it here (even though it
# seems irrelevant)
sim.compile(loss="mse", metrics=["accuracy"])
# assert layer.model.layers[0] in sim.tensor_graph.layers
assert layer.model.weights[0] in sim.keras_model.weights
def test_extra_feeds(Simulator):
# set up a tensornode that will fail unless a value is fed in for the
# placeholder
class NodeFunc(object):
def pre_build(self, *_):
self.ph = tf.placeholder_with_default(False, ())
def __call__(self, t, x):
with tf.device("/cpu:0"):
check = tf.Assert(self.ph, [t])
with tf.control_dependencies([check]):
y = tf.identity(x)
return y
with nengo.Network() as net:
a = nengo.Node([0])
b = TensorNode(NodeFunc(), size_in=1, size_out=1)
nengo.Connection(a, b)
p = nengo.Probe(b)
with Simulator(net) as sim:
with pytest.raises(tf.errors.InvalidArgumentError):
sim.run_steps(10)
sim.run_steps(10, extra_feeds={b.tensor_func.ph: True})
inputs = {a: np.zeros((1, 10, 1))}
targets = {p: np.zeros((1, 10, 1))}
with pytest.raises(tf.errors.InvalidArgumentError):
sim.train(inputs, targets, tf.train.GradientDescentOptimizer(1))
sim.train(inputs,
targets,
tf.train.GradientDescentOptimizer(1),
extra_feeds={b.tensor_func.ph: True})
with pytest.raises(tf.errors.InvalidArgumentError):
sim.loss(inputs, targets, "mse")
sim.loss(inputs, targets, "mse", extra_feeds={b.tensor_func.ph: True})
def test_validation(Simulator):
with nengo.Network() as net:
# not a callable
with pytest.raises(ValidationError, match="function or Keras Layer"):
TensorNode([0])
# size out < 1
with pytest.raises(ValidationError, match="must be >= 1"):
TensorNode(lambda t: t, shape_out=(0, ))
# wrong call signature
with pytest.raises(ValidationError,
match="function produced an error"):
TensorNode(lambda a, b, c: a)
# wrong call signature
with pytest.raises(ValidationError,
match="signature produced an error"):
TensorNode(tf.keras.layers.Lambda(lambda a, b, c: a))
# returning None
with pytest.raises(ValidationError, match="must return a Tensor"):
TensorNode(lambda x: None, shape_in=(1, ), pass_time=False)
# returning non-tensor
with pytest.raises(ValidationError, match="must return a Tensor"):
TensorNode(lambda x: [0], shape_in=(1, ), pass_time=False)
# no input
with pytest.raises(ValidationError,
match="either shape_in or pass_time"):
TensorNode(None, pass_time=False)
# correct output
n = TensorNode(lambda t: tf.zeros((5, 2)))
assert n.size_out == 2
# can't run tensornode in regular Nengo simulator
with nengo.Simulator(net) as sim:
with pytest.raises(SimulationError,
match="Cannot call TensorNode output"):
sim.step()
# these tensornodes won't be validated at creation, because size_out
# is specified. instead the validation occurs when the network is built
# None output
with nengo.Network() as net:
TensorNode(lambda t: None, shape_out=(2, ))
with pytest.raises(ValidationError, match="must return a Tensor"):
Simulator(net)
# wrong number of dimensions
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((1, 2, 2)), shape_out=(2, ))
with pytest.raises(ValidationError, match="should have size"):
Simulator(net)
# wrong minibatch size
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2)), shape_out=(2, ))
with pytest.raises(ValidationError, match="should have batch size"):
Simulator(net, minibatch_size=2)
# wrong output d
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2)), shape_out=(3, ))
with pytest.raises(ValidationError, match="should have size"):
Simulator(net, minibatch_size=3)
# wrong dtype
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2), dtype=tf.int32), shape_out=(2, ))
with pytest.raises(ValidationError, match="should have dtype"):
Simulator(net, minibatch_size=3)
# make sure that correct output _does_ pass
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2), dtype=t.dtype), shape_out=(2, ))
with Simulator(net, minibatch_size=3):
pass
def test_validation(Simulator):
with nengo.Network() as net:
# not a callable
with pytest.raises(ValidationError):
TensorNode([0])
# size out < 1
with pytest.raises(ValidationError):
TensorNode(lambda t: t, size_out=0)
# wrong call signature
with pytest.raises(ValidationError):
TensorNode(lambda a, b, c: a)
# returning None
with pytest.raises(ValidationError):
TensorNode(lambda x: None)
# returning non-tensor
with pytest.raises(ValidationError):
TensorNode(lambda x: [0])
# returning wrong number of dimensions
with pytest.raises(ValidationError):
TensorNode(lambda t: tf.zeros((2, 2, 2)))
# correct output
n = TensorNode(lambda t: tf.zeros((5, 2)))
assert n.size_out == 2
# can't run tensornode in regular Nengo simulator
with nengo.Simulator(net) as sim:
with pytest.raises(SimulationError):
sim.step()
# these tensornodes won't be validated at creation, because size_out
# is specified. instead the validation occurs when the network is built
# None output
with nengo.Network() as net:
TensorNode(lambda t: None, size_out=2)
with pytest.raises(ValidationError):
Simulator(net)
# wrong number of dimensions
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((1, 2, 2)), size_out=2)
with pytest.raises(ValidationError):
Simulator(net)
# wrong minibatch size
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2)), size_out=2)
with pytest.raises(ValidationError):
Simulator(net, minibatch_size=2)
# wrong output d
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2)), size_out=3)
with pytest.raises(ValidationError):
Simulator(net, minibatch_size=3)
# wrong dtype
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2), dtype=tf.int32), size_out=2)
with pytest.raises(ValidationError):
Simulator(net, minibatch_size=3)
# make sure that correct output _does_ pass
with nengo.Network() as net:
TensorNode(lambda t: tf.zeros((3, 2), dtype=t.dtype), size_out=2)
with Simulator(net, minibatch_size=3):
pass
def __init__(self, input_size):
with self:
x = nengo.Node(input_size)
x = TensorNode(Conv2DNode)
