train_tools.py 文件源码

def minimize_clipped(optimizer, loss, clip_value, return_gvs=False, soft=False, **kwargs):
    """Computes a train_op with clipped gradients in the range [-clip_value, clip_value]

    :param optimizer: Tensorflow optimizer object
    :param loss: tensor
    :param clip_value: scalar value
    :param return_gvs: returns list of tuples of (gradient, parameter) for trainable variables
    :param kwargs: kwargs for optimizer.compute_gradients function
    :return: train_step
    """

    gvs = optimizer.compute_gradients(loss, **kwargs)
    clipped_gvs = [(g, v) for (g, v) in gvs if g is not None]

    if not soft:
        clipped_gvs = [(tf.clip_by_value(g, -clip_value, clip_value), v) for (g, v) in clipped_gvs]

    else:
        n_elems = 0
        norm_squared = 0.
        for g, v in gvs:
            n_elems += tf.reduce_prod(tf.shape(g))
            norm_squared += tf.reduce_sum(g ** 2)

        norm_squared /= tf.to_float(n_elems)
        inv_norm = gen_math_ops.rsqrt(norm_squared)
        cond = tf.greater(norm_squared, clip_value ** 2)

        def clip(x):
            return tf.cond(cond, lambda: clip_value * x * inv_norm, lambda: x)

        clipped_gvs = [(clip(g), v) for (g, v) in clipped_gvs]

    train_step = optimizer.apply_gradients(clipped_gvs)

    if return_gvs:
        train_step = (train_step, gvs)
    return train_step