Question

我想定义一个自定义LearningRateSchedule，但是AutoGraph似乎很难转换它。以下代码在没有@ tf.function的情况下可以正常工作。但是使用@tf.function

时会引发错误

def linear_interpolation(l, r, alpha):
    return l + alpha * (r - l)

class TFPiecewiseSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    # This class currently cannot be used in @tf.function, 
    # Since tf.cond See the following link for details
    def __init__(self, endpoints, end_learning_rate=None, name=None):
        """Piecewise schedule.
        endpoints: [(int, int)]
            list of pairs `(time, value)` meanining that schedule should output
            `value` when `t==time`. All the values for time must be sorted in
            an increasing order. When t is between two times, e.g. `(time_a, value_a)`
            and `(time_b, value_b)`, such that `time_a <= t < time_b` then value outputs
            `interpolation(value_a, value_b, alpha)` where alpha is a fraction of
            time passed between `time_a` and `time_b` for time `t`.
        outside_value: float
            if the value is requested outside of all the intervals sepecified in
            `endpoints` this value is returned. If None then AssertionError is
            raised when outside value is requested.
        """
        super().__init__()
        idxes = [e[0] for e in endpoints]
        assert idxes == sorted(idxes)
        self.end_learning_rate = end_learning_rate or endpoints[-1][1]
        self.endpoints = endpoints
        self.name=name

    def __call__(self, step):
        if step < self.endpoints[0][0]:
            return self.endpoints[0][1]
        else:
            for (l_t, l), (r_t, r) in zip(self.endpoints[:-1], self.endpoints[1:]):
                if l_t <= step < r_t:
                    alpha = float(step - l_t) / (r_t - l_t)
                    return linear_interpolation(l, r, alpha)

        # t does not belong to any of the pieces, so doom.
        assert self.end_learning_rate is not None
        return self.end_learning_rate

    def get_config(self):
        return dict(
            endpoints=self.endpoints,
            end_learning_rate=self.end_learning_rate,
            name=self._name,
        )

lr = TFPiecewiseSchedule([[10, 1e-3], [20, 1e-4]])
@tf.function
def f(x):
    l = layers.Dense(10)
    with tf.GradientTape() as tape:
        y = l(x)
        loss = tf.reduce_mean(y**2)

    grads = tape.gradient(loss, l.trainable_variables)
    opt = tf.keras.optimizers.Adam(lr)
    opt.apply_gradients(zip(grads, l.trainable_variables))

f(tf.random.normal((2, 3)))

错误消息显示：

：10楼*           opt.apply_gradients（zip（grads，l.trainable_variables））

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:437 apply_gradients           apply_state = self._prepare（var_list）

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:614 _prepare           self._prepare_local（var_device，var_dtype，apply_state）

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/adam.py:154 _prepare_local           超级（亚当，自我）。_prepare_local（var_device，var_dtype，apply_state）

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:620 _prepare_local           lr_t = array_ops.identity（self._decayed_lr（var_dtype））

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:672 _decayed_lr           lr_t = math_ops.cast（lr_t（local_step），var_dtype）

：32 致电           如果步骤
/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:765 bool           self._disallow_bool_casting（）

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:531 _disallow_bool_casting           “使用tf.Tensor作为Python bool”）

/Users/aptx4869/anaconda3/envs/drl/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:518 _disallow_when_autograph_enabled           “使用@ tf.function直接修饰它。”。format（task））

OperatorNotAllowedInGraphError：不允许将tf.Tensor作为Python bool使用：AutoGraph没有转换此函数。尝试直接使用@ tf.function装饰它。

我认为该错误是由于if语句引起的，所以我用以下代码替换了__call__函数的内容。但是会出现几乎相同的错误。

        def compute_lr(step):
            for (l_t, l), (r_t, r) in zip(self.endpoints[:-1], self.endpoints[1:]):
                if l_t <= step < r_t:
                    alpha = float(step - l_t) / (r_t - l_t)
                    return linear_interpolation(l, r, alpha)

            # t does not belong to any of the pieces, so doom.
            assert self.end_learning_rate is not None
            return self.end_learning_rate
        return tf.cond(tf.less(step, self.endpoints[0][0]), lambda: self.endpoints[0][1], lambda: compute_lr(step))

我应该怎么做才能使代码按我的意愿工作？

Answer 1

markdown格式化程序将错误消息弄乱了，但是似乎aa.cocnumber函数本身未由AutoGraph处理。在错误消息中，转换后的功能标有星号。这是Adam优化器中的错误。无论如何，您可以直接使用tf对其进行注释。函数将被拾取：

__call__

也就是说，代码中有一些AutoGraph不喜欢的东西：@tf.function def __call__(self, step):，从循环返回，链不等式-尽可能使用基本结构更安全。令人遗憾的是，您仍然会收到错误，这令人困惑。像这样重写它应该起作用：

zip

最后一个问题-当事物创建变量时@tf.function def __call__(self, step): if step < self.endpoints[0][0]: return self.endpoints[0][1] else: # Can't return from a loop lr = self.end_learning_rate # Since it needs to break based on the value of a tensor, loop # needs to be a tf.while_loop for pair in tf.stack([self.endpoints[:-1], self.endpoints[1:]], axis=1): left, right = tf.unstack(pair) l_t, l = tf.unstack(left) r_t, r = tf.unstack(right) # Chained inequalities not supported yet if l_t <= step and step < r_t: alpha = float(step - l_t) / (r_t - l_t) lr = linear_interpolation(l, r, alpha) break return lr不喜欢它，因此您需要将层和优化器的创建移到外部：

tf.function

我希望这会有所帮助！

@ tf.function中的if-else

1 个答案: