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)