Search code examples
tensorflowgpflow

optimization in gpflow 2: Why set autograph=False?

in the current notebook tutorials (gpflow 2.0), all @tf.function tags include the option autograph=False, e.g. (https://gpflow.readthedocs.io/en/2.0.0-rc1/notebooks/advanced/gps_for_big_data.html):

@tf.function(autograph=False)
def optimization_step(optimizer, model: gpflow.models.SVGP, batch):
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(model.trainable_variables)
        objective = - model.elbo(*batch)
        grads = tape.gradient(objective, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return objective

Does anyone know why that is the case, or what the reasoning behind this is? As far as I understood, autograph=True simply allows for python control flow to be translated to a graph structure. Does setting/leaving it to true, even if the functionality is not required, have any drawbacks?

My guess would have been that its just a small overhead at compile time of the graph, but should be negligible. Is that wrong?

Thanks

Solution

  • The reason we set autograph to False in most of the tf.function wrapped objectives is because GPflow makes use a multi-dispatch Dispatcher which internally uses generators. TensorFlow, however, can not deal with generator objects in autograph mode (see Capabilities and Limitations of AutoGraph), which leads to these warning:

    WARNING:tensorflow:Entity <bound method Dispatcher.dispatch_iter of <dispatched sample_conditional>> appears to be a generator function. It will not be converted by AutoGraph.
WARNING: Entity <bound method Dispatcher.dispatch_iter of <dispatched sample_conditional>> appears to be a generator function. It will not be converted by AutoGraph.
WARNING:tensorflow:Entity <bound method Dispatcher.dispatch_iter of <dispatched conditional>> appears to be a generator function. It will not be converted by AutoGraph.
WARNING: Entity <bound method Dispatcher.dispatch_iter of <dispatched conditional>> appears to be a generator function. It will not be converted by AutoGraph.

    We've know about this issue for a while but haven't got around to actually fixing it - thanks for bringing this back to our attention. I've just created a PR which fixes this issue and does not require you to set autograph to False anymore. I expect this PR to be merged fairly soon.