Deep copying a tf.function-decorated function?

user_1900954 0
Viewed 21

Some code I'm using is trying but failing to copy.deepcopy a @tf.function-decorated function (as part of pickling). Is there a correct way to do this?

Simplified reproduction

Define any simple @tf.function-decorated function:

>>> import tensorflow as tf
>>> @tf.function
... def foo(x): return x * 2
>>> foo
<tensorflow.python.eager.def_function.Function object at 0x7fc98bec0950>

Deepcopying it before initialisation works fine:

>>> import copy
>>> copy.deepcopy(foo)
<tensorflow.python.eager.def_function.Function object at 0x7fc98bec8b50>

However, it fails after initialisation:

>>> foo(tf.constant([3.]))
<tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>
>>> copy.deepcopy(foo)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/.pyenv/versions/3.7.13/lib/python3.7/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  ....
  File "/home/user/.pyenv/versions/3.7.13/lib/python3.7/copy.py", line 169, in deepcopy
    rv = reductor(4)
TypeError: can't pickle _thread.RLock objects

Comments

  1. I don't mind if the copied version is uninitialized and so has to be recompiled on first use, but I don't want to have to recompile the original function every time I copy it.
  2. The tensorflow.python.eager.def_function.Function class has a __getstate__ method that omits the _lock attribute and other "unpickleable objects". However, while this might permit shallow copies it doesn't handle deep ones, since the class objects also contain attributes of the following types, which themselves contain locks somewhere inside them:
_stateful_fn: <class 'tensorflow.python.eager.function.Function'>
_stateless_fn: <class 'tensorflow.python.eager.function.Function'>
_lifted_initializer_graph: <class 'tensorflow.python.framework.func_graph.FuncGraph'>
_graph_deleter: <class 'tensorflow.python.eager.def_function.FunctionDeleter'>
_concrete_stateful_fn: <class 'tensorflow.python.eager.function.ConcreteFunction'>

Environment

Tested in Python 3.7.13 on Ubuntu with both TF 2.8 and TF 2.4.

tensorflow2.0 tensorflow python
1 Answers

Figured out one possible solution just after posting.

Putting the function inside a class seems to let it be copied:

>>> class Foo:
...    @tf.function
...    def __call__(self, x): return x * 2
>>> foo = Foo()
>>> foo(tf.constant([3.]))
<tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>
>>> foo_copy = copy.deepcopy(foo)
>>> foo_copy(tf.constant([3.]))
<tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>

Interestingy, the copied function does indeed seem to be uninitialised (but that's ok for me).

>>> foo.__call__._get_tracing_count()
1
>>> copy.deepcopy(foo).__call__._get_tracing_count()
0
user_19009540
Related