如何用 tf.function 装饰一个以 tf.variable 作为参数的函数，最重要的是使用输入签名

Question

我遇到了一个问题，我需要修改 Tensorflow 函数内的一个变量。那我需要将此函数转换为张量流图。问题是变量的大小不固定。示例：它可以是形状为 (3,) 或 (2,) 的 tenosr。这就是函数将这个变量作为参数的原因，以便它可以修改它并return它。

这里是一个 class 的例子，它包含一个函数 call，这个函数有两个参数 (x,v)。 x 是 Tf.tensor，v 是 tf.Variable。 v 被赋值为 x*v 的乘积。

import tensorflow as tf

class MyModule(tf.Module):
  def __init__(self):
    pass

  @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.int32), tf.TensorSpec(shape=[None], dtype=tf.int32)])
  def __call__(self, x, v):
    v.assign(x*v, read_value=False)
    return v

tf.config.run_functions_eagerly(False)
x = tf.constant([10,10])
v = tf.Variable(2*tf.ones_like(x), trainable=False)

module = MyModule()
module(x, v)

这在 eager 模式下按预期工作，但在图形模式下我收到以下错误： AttributeError: 'Tensor' 对象没有属性 'assign'

我知道是因为tf.Variable的签名。我的问题是，如果当前签名产生错误，我如何指定 tf.Variable 的签名？

Answer 1

其实有一种操作可以达到你想要的效果，但是public API中没有列出。请注意，这可能不是最佳做法。

您需要 resource_variable_ops，可以在 tensorflow.python.ops 下找到。

import tensorflow as tf
from tensorflow.python.ops import resource_variable_ops

class MyModule(tf.Module):
  def __init__(self):
    pass

  @tf.function(input_signature=[
                                tf.TensorSpec(shape=[None], dtype=tf.int32), 
                                resource_variable_ops.VariableSpec(shape=[None], dtype=tf.int32)
                                ])
  def __call__(self, x, v):
    v.assign(x*v, read_value=False)
    return v

x = tf.constant([10,10])
v = tf.Variable(2*tf.ones_like(x), trainable=False)

module = MyModule()
module(x, v)

Answer 2

用简单的方法就可以完成

[样本]:

import tensorflow as tf

class MyModule(tf.Module):
    def __init__(self, v):
        self.v = v
        pass

    @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.int32), tf.TensorSpec(shape=None, dtype=tf.int32)])
    def __call__(self, x, v):
        self.v.assign( x * v, read_value=False )
        return self.v

x = tf.constant( tf.random.uniform(shape=[2,1], maxval=3, dtype=tf.int32) )
v = tf.Variable([[1], [2]])
module = MyModule(v)

print( module(x, v) )

#############################################
x = tf.constant( tf.random.uniform(shape=[3,1], maxval=3, dtype=tf.int32) )
v = tf.Variable([[1], [2], [3]])
module = MyModule(v)

print( module(x, v) )

[输出]:

tf.Tensor(
[[1]
 [0]], shape=(2, 1), dtype=int32)
tf.Tensor(
[[2]
 [0]
 [6]], shape=(3, 1), dtype=int32)

Answer 3

这是我找到的解决方案：

class MyModule(tf.Module):
    def __init__(self):
        self.v = tf.Variable([[]], shape=tf.TensorShape([None]*2), dtype=tf.int32)

    @tf.function(input_signature=[tf.TensorSpec(shape=[None]*2, dtype=tf.int32), tf.TensorSpec(shape=[None]*2, dtype=tf.int32)])
    def __call__(self, x, v):
      self.v.assign( x * v, read_value=False)
      return self.v

在模型的构造函数中用空值定义 tf.Variable 解决了问题。

如何用 tf.function 装饰一个以 tf.variable 作为参数的函数，最重要的是使用输入签名

How to decorate a function that takes a tf.variable as a parameter with tf.function and most importantly using input signature

deep-learning

tensorflow