时间戳

首页 > 解决方案 > 当 tensorflow 2.0 中的数据没有变化时，如何为梯度更新构建循环？

tensorflow2.0 - 当 tensorflow 2.0 中的数据没有变化时，如何为梯度更新构建循环？

问题描述

这两个代码片段显示了使用相同的数据进行 n 次更新，其中一个使用持久梯度磁带，另一个只是一遍又一遍地调用它。性能差异似乎约为 2 倍。有没有更好的方法来构建这个未来？我想将数据移动到设备对 GPU 很重要？

@tf.function
def train_n_steps_same_data(tau, y, n=1):
    """
    In [218]: %timeit r = q.train_n_steps_same_data(q.tau, q.y, n=100)
    25.3 ms ± 926 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
    """
    with tf.GradientTape(persistent=True) as tape:
        d = model([tau, y])
        loss = tf.reduce_mean(d['rho'])
    for i in range(n):
        gradients = tape.gradient(loss, model.trainable_variables)
        l = optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    names = [x.name for x in gradients]
    g = dict(zip(names, gradients))
    reduced = dict()
    reduced['loss'] = loss
    return reduced, d, g

@tf.function
def train_n_steps_same_data2(tau, y, n=1):
    """
    In [220]: %timeit r = q.train_n_steps_same_data2(q.tau, q.y, n=100)
    41.6 ms ± 1.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
    """
    for i in range(n):
        with tf.GradientTape() as tape:
            d = model([tau, y])
            loss = tf.reduce_mean(d['rho'])
        gradients = tape.gradient(loss, model.trainable_variables)
        l = optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    names = [x.name for x in gradients]
    g = dict(zip(names, gradients))
    reduced = dict()
    reduced['loss'] = loss
    return reduced, d, g

标签： tensorflow2.0

解决方案

第一种方法肯定更好。您创建一个磁带对象并在循环内重复使用它。相反，第二个函数在每次迭代时创建和销毁磁带对象。

但是，您在第一个训练循环中遗漏了一个非常重要的部分：您的磁带是持久的。因此，在您使用它之后，您必须使用 手动删除它del tape，否则会导致内存泄漏。

range另一个建议是当你用 装饰函数时不要使用tf.function，而是使用tf.range（一般情况下，尽可能使用tf.*与 Python 构造等效的方法，请参阅这篇文章）

推荐阅读