Using TensorFlow autograph in NengoDL

Backends
1

Can I use autograph while building the TensorFlow computation graph implementing some NengoDL OpBuilder's build_step() function?

For example, if I want conditional execution of part of my code, am I limited to using tf.cond() to insert a decision node or can I use Python declarative instructions and convert using @tf.function?

For example, will this traditional TensorFlow code:

@Builder.register( MyLearningRule )
class MyLearningRuleBuilder( OpBuilder ):
    ...
    def build_step(self, signals):
        def if_error_over_threshold():
            # do something
        tf.cond( tf.reduce_any( tf.greater( tf.abs( local_error ), error_threshold ) ),
        true_fn=if_error_over_threshold,
        false_fn=lambda: None )

be equivalent in NengoDL to this one using autograph:

@Builder.register( MyLearningRule )
class MyLearningRuleBuilder( OpBuilder ):
    ...
    @tf.function
    def build_step(self, signals):
        if tf.reduce_any( tf.greater( tf.abs( local_error ), error_threshold ) ):
            # do something
        else:
            # do nothing

?

2

Hi Tioz,

Welcome to the Nengo forum.

Currently in NengoDL we do all the TensorFlow building in graph mode, not eager mode, so the autograph/tf.function functionality won’t work. We do plan on switching to eager mode soon, now that the performance issues have been largely resolved, at which point all that autograph stuff should work.

3

Excellent, I’m looking forwards to it!

4

As NengoDL 3.3.0 has added support for TensorFlow 2.0 syntax, I would appreciate an update on my original question: how can I port my custom learning rule from 1.0 to 2.0 syntax? :slight_smile: Is it just a case of substituting my tf.cond() statements with a Python if?

5

Autograph is still disabled within the NengoDL build process (see https://github.com/nengo/nengo-dl/blob/master/nengo_dl/tensor_graph.py#L396), because it causes occasional errors (and in general I prefer the explicitness of using tf.cond, rather than the somewhat unknown code manipulation of AutoGraph). However, you could definitely enable autograph and experiment with it, it might just work for your use case out of the box.

6

My code works fine as it is so I guess I’ll wait for the feature to officially be supported before experimenting with it :slight_smile: