Re-enable eager mode

.appveyor.yml

@@ -16,7 +16,7 @@ environment:
       TF_VERSION: tensorflow --pre
       PYTHON_VERSION: C:\Python38-x64
     - NENGO_VERSION: nengo[tests]==3.0.0
-      TF_VERSION: tensorflow==2.0.0
+      TF_VERSION: tensorflow==2.2.0
       NUMPY_VERSION: numpy==1.16.0
       PYTHON_VERSION: C:\Python35-x64
 

.nengobones.yml

@@ -204,9 +204,11 @@ travis_yml:
     - script: test-coverage
       env:
         NENGO_VERSION: nengo[tests]==3.0.0
-        TF_VERSION: tensorflow==2.0.0
+        TF_VERSION: tensorflow==2.2.0
         NUMPY_VERSION: numpy==1.16.0
       python: 3.5
+    - script: test-coverage
+      test_args: --graph-mode
     - stage: advanced
       script: remote-docs
       env:
@@ -216,8 +218,7 @@ travis_yml:
       env:
         TF_VERSION: tensorflow-gpu
         GPU_NUM: 0
-    - test_args: --unroll-simulation 5 --simulator-only
-    - test_args: --dtype float64 --simulator-only
+    - test_args: --unroll-simulation 5 --dtype float64 --simulator-only
     - test_args: --inference-only --simulator-only
   pypi_user: drasmuss
   deploy_dists:

.templates/setup.py.template

@@ -43,7 +43,7 @@ else:
 install_req = [
     "nengo>=3.0.0",
     "numpy>=1.16.0",
-    "%s>=2.0.0" % tf_req,
+    "%s>=2.2.0" % tf_req,
     "jinja2>=2.10.1",
     "packaging>=20.0",
     "progressbar2>=3.39.0",

.travis.yml

@@ -52,10 +52,14 @@ jobs:
   -
     env:
       NENGO_VERSION="nengo[tests]==3.0.0"
-      TF_VERSION="tensorflow==2.0.0"
+      TF_VERSION="tensorflow==2.2.0"
       NUMPY_VERSION="numpy==1.16.0"
       SCRIPT="test-coverage"
     python: 3.5
+  -
+    env:
+      SCRIPT="test-coverage"
+      TEST_ARGS="--graph-mode"
   -
     env:
       TF_VERSION="tensorflow-gpu"
@@ -69,10 +73,7 @@ jobs:
       SCRIPT="remote-examples"
   -
     env:
-      TEST_ARGS="--unroll-simulation 5 --simulator-only"
-  -
-    env:
-      TEST_ARGS="--dtype float64 --simulator-only"
+      TEST_ARGS="--unroll-simulation 5 --dtype float64 --simulator-only"
   -
     env:
       TEST_ARGS="--inference-only --simulator-only"

CHANGES.rst

@@ -18,7 +18,7 @@ Release history
    - Deprecated
    - Removed
 
-3.2.1 (unreleased)
+3.3.0 (unreleased)
 ------------------
 
 **Added**
@@ -27,6 +27,12 @@ Release history
 - Compatible with TensorFlow 2.3.0. (`#159`_)
 - Added support for ``nengo.Tanh``, ``nengo.RegularSpiking``,
   ``nengo.StochasticSpiking``, and ``nengo.PoissonSpiking`` neuron types. (`#159`_)
+- Added ``nengo_dl.configure_settings(learning_phase=True/False)`` configuration
+  option. This mimics the previous behaviour of
+  ``tf.keras.backend.learning_phase_scope`` (which was deprecated by TensorFlow). That
+  is, if you would like to override the default behaviour so that, e.g., ``sim.predict``
+  runs in training mode, set ``nengo_dl.configure_settings(learning_phase=True)``.
+  (`#163`_)
 
 **Changed**
 
@@ -37,6 +43,29 @@ Release history
 - A warning will now be raised if activation types are passed to
   ``Converter.swap_activations`` that aren't actually in the model. (`#168`_)
 - Updated TensorFlow installation instruction in documentation. (`#170`_)
+- NengoDL will now use TensorFlow's eager mode by default. The previous graph-mode
+  behaviour can be restored by calling ``tf.compat.v1.disable_eager_execution()``, but
+  we cannot guarantee that that behaviour will be supported in the future. (`#163`_)
+- NengoDL will now use TensorFlow's "control flow v2" by default. The previous
+  behaviour can be restored by calling ``tf.compat.v1.disable_control_flow_v2()``, but
+  we cannot guarantee that that behaviour will be supported in the future. (`#163`_)
+- NengoDL will now default to allowing TensorFlow's "soft placement" logic, meaning
+  that even if you specify an explicit device like ``"/gpu:0"``, TensorFlow may not
+  allocate an op to that device if there isn't a compatible implementation available.
+  The previous behaviour can be restored by calling
+  ``tf.config.set_soft_device_placement(False)``. (`#163`_)
+- Internal NengoDL ``OpBuilder`` classes now separate the "pre build" stage from
+  ``OpBuilder.__init__`` (so that the same ``OpBuilder`` class can be re-used across
+  multiple ``calls``, rather than instantiating a new ``OpBuilder`` each time). Note
+  that this has no impact on front-end users, this is
+  only relevant to anyone that has implemented a custom build class. The
+  logic that would previously have gone in ``OpBuilder.__init__`` should now go in
+  ``OpBuilder.build_pre``. In addition, the ``ops`` argument has been removed
+  from ``OpBuilder.build_pre``; that will be passed to ``OpBuilder.__init__`` (
+  and will be available in ``build_pre`` as ``self.ops``). Similarly, the ``ops`` and
+  ``config`` argument have been removed from ``build_post``, and can instead be
+  accessed through ``self.ops/config``. (`#163`_)
+- Minimum TensorFlow version is now 2.2.0. (`#163`_)
 
 **Fixed**
 
@@ -50,11 +79,16 @@ Release history
 - Fix bug when probing slices of certain probeable attributes (those that are
   directly targeting a Signal in the model). (`#164`_)
 
+**Removed**
+
+- Removed ``nengo_dl.utils.print_op`` (use ``tf.print`` instead). (`#163`_)
+
 .. _#149: https://github.com/nengo/nengo-dl/pull/149
 .. _#151: https://github.com/nengo/nengo-dl/pull/151
 .. _#153: https://github.com/nengo/nengo-dl/pull/153
 .. _#159: https://github.com/nengo/nengo-dl/pull/159
 .. _#161: https://github.com/nengo/nengo-dl/pull/161
+.. _#163: https://github.com/nengo/nengo-dl/pull/163
 .. _#164: https://github.com/nengo/nengo-dl/pull/164
 .. _#168: https://github.com/nengo/nengo-dl/pull/168
 .. _#170: https://github.com/nengo/nengo-dl/pull/170

conftest.py

@@ -12,9 +12,19 @@ def pytest_configure(config):
     config.addinivalue_line(
         "markers", "performance: mark tests that depend on a specific benchmark machine"
     )
+    config.addinivalue_line(
+        "markers", "eager-only: mark tests that only work in eager mode"
+    )
+
+    if config.getvalue("--graph-mode"):
+        tf.compat.v1.disable_eager_execution()
+        tf.compat.v1.disable_control_flow_v2()
 
 
 def pytest_runtest_setup(item):
+    # NOTE: this hook will not be called when running --pyargs nengo, so don't
+    # put anything here that we want to run during those tests
+
     if item.get_closest_marker("gpu", False) and not utils.tf_gpu_installed:
         pytest.skip("This test requires tensorflow-gpu")
     elif (
@@ -31,8 +41,10 @@ def pytest_runtest_setup(item):
         "--performance"
     ):
         pytest.skip("Skipping performance test")
-
-    tf.keras.backend.clear_session()
+    elif item.get_closest_marker("eager_only", False) and item.config.getvalue(
+        "--graph-mode"
+    ):
+        pytest.skip("Skipping eager-only test")
 
 
 def pytest_addoption(parser):
@@ -67,6 +79,12 @@ def pytest_addoption(parser):
         default=False,
         help="Run performance tests",
     )
+    parser.addoption(
+        "--graph-mode",
+        action="store_true",
+        default=False,
+        help="Run tests in graph (not eager) mode",
+    )
 
 
 @pytest.fixture(scope="session")
@@ -76,3 +94,11 @@ def Simulator(request):
     """
 
     return make_test_sim(request)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def clear_session(request):
+    # free up resources between tests.
+    # we do this here, rather than in e.g. pytest_runtest_setup, because we want this
+    # to take effect when running the nengo core tests as well
+    tf.keras.backend.clear_session()

docs/examples/from-nengo.ipynb

@@ -39,6 +39,7 @@
     "\n",
     "\n",
     "warnings.simplefilter(\"ignore\")\n",
+    "tf.get_logger().addFilter(lambda rec: \"Tracing is expensive\" not in rec.msg)\n",
     "\n",
     "# we'll control the random seed in this example to make sure things stay\n",
     "# consistent, but the results don't depend significantly on the seed\n",

docs/examples/from-tensorflow.ipynb

@@ -30,7 +30,8 @@
     "import nengo_dl\n",
     "\n",
     "\n",
-    "warnings.simplefilter(\"ignore\")"
+    "warnings.simplefilter(\"ignore\")\n",
+    "tf.get_logger().addFilter(lambda rec: \"Tracing is expensive\" not in rec.msg)"
    ]
   },
   {

docs/examples/keras-to-snn.ipynb

@@ -198,9 +198,13 @@
     "    nengo_input = nengo_converter.inputs[inp]\n",
     "    nengo_output = nengo_converter.outputs[dense]\n",
     "    \n",
-    "    # add a probe to the first convolutional layer to record activity\n",
+    "    # add a probe to the first convolutional layer to record activity.\n",
+    "    # we'll only record from a subset of neurons, to save memory.\n",
+    "    sample_neurons = np.linspace(\n",
+    "        0, np.prod(conv0.shape[1:]), 1000, endpoint=False, dtype=np.int32,\n",
+    "    )\n",
     "    with nengo_converter.net:\n",
-    "        conv0_probe = nengo.Probe(nengo_converter.layers[conv0])\n",
+    "        conv0_probe = nengo.Probe(nengo_converter.layers[conv0][sample_neurons])\n",
     "            \n",
     "    # repeat inputs for some number of timesteps\n",
     "    tiled_test_images = np.tile(test_images[:n_test], (1, n_steps, 1))\n",
@@ -232,14 +236,7 @@
     "        plt.axis('off')\n",
     "        \n",
     "        plt.subplot(1, 3, 2)\n",
-    "        sample_neurons = np.linspace(\n",
-    "            0,\n",
-    "            data[conv0_probe].shape[-1], \n",
-    "            1000,\n",
-    "            endpoint=False,\n",
-    "            dtype=np.int32,\n",
-    "        )\n",
-    "        scaled_data = data[conv0_probe][ii, :, sample_neurons].T * scale_firing_rates\n",
+    "        scaled_data = data[conv0_probe][ii] * scale_firing_rates\n",
     "        if isinstance(activation, nengo.SpikingRectifiedLinear):\n",
     "            scaled_data *= 0.001\n",
     "            rates = np.sum(scaled_data, axis=0) / (n_steps * nengo_sim.dt)\n",

docs/examples/tensorflow-models.ipynb

@@ -19,8 +19,6 @@
    "source": [
     "%matplotlib inline\n",
     "\n",
-    "import threading\n",
-    "\n",
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
     "import tensorflow as tf\n",
@@ -201,13 +199,6 @@
     "        # load the weights we saved above\n",
     "        self.model.load_weights(model_weights)\n",
     "        \n",
-    "        # this part is just for compatibility between different\n",
-    "        # TensorFlow versions\n",
-    "        if not hasattr(self.model, \"_metrics_lock\"):\n",
-    "            self.model._metrics_lock = threading.Lock()\n",
-    "            for lay in self.model.layers:\n",
-    "                lay._metrics_lock = threading.Lock()\n",
-    "        \n",
     "    def call(self, inputs):\n",
     "        # apply the model to the inputs\n",
     "        return self.model(inputs)"

docs/migration-guide.rst

@@ -98,6 +98,11 @@ Simulator changes
             sim.evaluate(
                 x={node: np.zeros((1, 1, 1))}, y={probe: np.zeros((1, 1, 1))})
 
+    .. testoutput::
+        :hide:
+
+        ...
+
 3.  **Extra simulator steps will no longer be hidden**.  When simulating a number of
     timesteps that is not evenly divisible by ``Simulator.unroll_simulation``,
     extra simulation steps will be executed (this is true in both 2 and 3).
@@ -169,8 +174,8 @@ Simulator changes
     equivalents during training.  However, sometimes it is useful to manually enable
     this swapping in other functions (for example, in order to evaluate the loss
     function on test data but with the swapped rate neuron models).  There were a couple
-    ways to do this in NengoDL 2; in NengoDL 3 it is all controlled through the Keras
-    ``learning_phase``.
+    ways to do this in NengoDL 2; in NengoDL 3 it is all controlled through the
+    ``learning_phase`` configuration option.
 
     NengoDL 2:
 
@@ -189,13 +194,20 @@ Simulator changes
 
     .. testcode::
 
-        with tf.keras.backend.learning_phase_scope(1):
-            with nengo_dl.Simulator(example_net) as sim:
-                sim.compile(loss=tf.losses.mse)
-                sim.evaluate(
-                    x={node: np.zeros((1, 1, 1))}, y={probe: np.zeros((1, 1, 1))})
+        with example_net:
+            nengo_dl.configure_settings(learning_phase=True)
+
+        with nengo_dl.Simulator(example_net) as sim:
+            sim.compile(loss=tf.losses.mse)
+            sim.evaluate(
+                x={node: np.zeros((1, 1, 1))}, y={probe: np.zeros((1, 1, 1))})
+
+            sim.run(1.0)
 
-                sim.run(1.0)
+    .. testoutput::
+        :hide:
+
+        ...
 
 6.  **TensorBoard functionality replaced by Keras TensorBoard callback**.  NengoDL
     allows data about training metrics or model parameters to be output and displayed in

docs/tips.rst

@@ -52,6 +52,17 @@ search depth by setting
 
 See :ref:`the documentation <config-planner>` for more details.
 
+TensorFlow reworked a lot of their internal implementation details in TensorFlow 2.0.
+But for some models, the pre-2.0 implementation will be faster. This behaviour can be
+restored by calling
+
+.. code-block:: python
+
+    tf.compat.v1.disable_eager_execution()
+    tf.compat.v1.disable_control_flow_v2()
+
+at the top of your script.
+
 Training a spiking deep network
 -------------------------------
 

nengo_dl/__init__.py

@@ -37,14 +37,6 @@
 
 del os
 
-import tensorflow as tf
-
-# disable control flow v2 for performance reasons
-# (see https://github.com/tensorflow/tensorflow/issues/33052)
-tf.compat.v1.disable_control_flow_v2()
-
-del tf
-
 # need to explicitly import these to trigger the builder registration
 from nengo_dl import (
     op_builders,

nengo_dl/benchmarks.py

@@ -11,7 +11,6 @@
 from nengo.utils.filter_design import cont2discrete
 import numpy as np
 import tensorflow as tf
-from tensorflow.python.eager import profiler
 
 import nengo_dl
 
@@ -219,7 +218,7 @@ def mnist(use_tensor_layer=True):
             #     rates = amplitude / (tau_ref + tau_rc * tf.log1p(1 / z))
             #     return rates
 
-            def mnist_node(x):  # pragma: no cover
+            def mnist_node(x):  # pragma: no cover (runs in TF)
                 x = tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation=nl)(x)
                 x = tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation=nl)(x)
                 x = tf.keras.layers.AveragePooling2D(pool_size=2, strides=2)(x)
@@ -677,12 +676,12 @@ def run_profile(
 
             for _ in range(reps):
                 if do_profile:
-                    profiler.start()
+                    tf.profiler.experimental.start("profile")
                 start = timeit.default_timer()
                 sim.fit(x, y, epochs=1, n_steps=n_steps)
                 exec_time = min(timeit.default_timer() - start, exec_time)
                 if do_profile:
-                    profiler.save("profile", profiler.stop())
+                    tf.profiler.experimental.stop()
 
         else:
             # run once to eliminate startup overhead
@@ -692,12 +691,12 @@ def run_profile(
 
             for _ in range(reps):
                 if do_profile:
-                    profiler.start()
+                    tf.profiler.experimental.start("profile")
                 start = timeit.default_timer()
                 sim.predict(x, n_steps=n_steps)
                 exec_time = min(timeit.default_timer() - start, exec_time)
                 if do_profile:
-                    profiler.save("profile", profiler.stop())
+                    tf.profiler.experimental.stop()
 
     exec_time /= n_batches