(x, b, axis, act, alpha, gain)
return func_nonzero_2nd_grad(x, b)
Fast CUDA implementation of `fused_bias_act()` using custom ops.
Fast CUDA implementation of `fused_bias_act()` using custom ops.
Fast CUDA implementation of fused_bias_act () using custom ops .
def _fused_bias_act_cuda(x, b, axis, act, alpha, gain):
    """Fast CUDA implementation of `fused_bias_act()` using custom ops."""

    # Validate arguments.
    x = tf.convert_to_tensor(x)
    empty_tensor = tf.constant([], dtype=x.dtype)
    b = tf.convert_to_tensor(b) if b is not None else empty_tensor
    act_spec = activation_funcs[act]
    assert b.shape.rank == 1 and (b.shape[0] == 0 or b.shape[0] == x.shape[axis])
    assert b.shape[0] == 0 or 0 <= axis < x.shape.rank
    if alpha is None:
        alpha = act_spec.def_alpha
    if gain is None:
        gain = act_spec.def_gain

    # Special cases.
    if act == 'linear' and b is None and gain == 1.0:
        return x
    if act_spec.cuda_idx is None:
        return _fused_bias_act_ref(x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)

    # CUDA kernel.
    cuda_kernel = _get_plugin().fused_bias_act
    cuda_kwargs = dict(axis=axis, act=act_spec.cuda_idx, alpha=alpha, gain=gain)

    # Forward pass: y = func(x, b).
    def func_y(x, b):
        y = cuda_kernel(x=x, b=b, ref=empty_tensor, grad=0, **cuda_kwargs)
        y.set_shape(x.shape)
        return y

    # Backward pass: dx, db = grad(dy, x, y)
    def grad_dx(dy, x, y):
        ref = {'x': x, 'y': y}[act_spec.ref]
        dx = cuda_kernel(x=dy, b=empty_tensor, ref=ref, grad=1, **cuda_kwargs)
        dx.set_shape(x.shape)
        return dx
    def grad_db(dx):
        if b.shape[0] == 0:
            return empty_tensor
        db = dx
        if axis < x.shape.rank - 1:
            db = tf.reduce_sum(db, list(range(axis + 1, x.shape.rank)))
        if axis > 0:
            db = tf.reduce_sum(db, list(range(axis)))
        db.set_shape(b.shape)
        return db

    # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y)
    def grad2_d_dy(d_dx, d_db, x, y):
        ref = {'x': x, 'y': y}[act_spec.ref]
        d_dy = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=1, **cuda_kwargs)
        d_dy.set_shape(x.shape)
        return d_dy
    def grad2_d_x(d_dx, d_db, x, y):
        ref = {'x': x, 'y': y}[act_spec.ref]
        d_x = cuda_kernel(x=d_dx, b=d_db, ref=ref, grad=2, **cuda_kwargs)
        d_x.set_shape(x.shape)
        return d_x

    # Fast version for piecewise-linear activation funcs.
    @tf.custom_gradient
    def func_zero_2nd_grad(x, b):
        y = func_y(x, b)
        @tf.custom_gradient
        def grad(dy):
            dx = grad_dx(dy, x, y)
            db = grad_db(dx)
            def grad2(d_dx, d_db):
                d_dy = grad2_d_dy(d_dx, d_db, x, y)
                return d_dy
            return (dx, db), grad2
        return y, grad

    # Slow version for general activation funcs.
    @tf.custom_gradient
    def func_nonzero_2nd_grad(x, b):
        y = func_y(x, b)
        def grad_wrap(dy):
            @tf.custom_gradient
            def grad_impl(dy, x):
                dx = grad_dx(dy, x, y)
                db = grad_db(dx)
                def grad2(d_dx, d_db):
                    d_dy = grad2_d_dy(d_dx, d_db, x, y)
                    d_x = grad2_d_x(d_dx, d_db, x, y)
                    return d_dy, d_x
                return (dx, db), grad2
            return grad_impl(dy, x)
        return y, grad_wrap

    # Which version to use?
    if act_spec.zero_2nd_grad:
        return func_zero_2nd_grad(x, b)
    return func_nonzero_2nd_grad(x, b)
def _fused_bias_act_cuda ( x , b , axis , act , alpha , gain ) : # Validate arguments. x = tf . convert_to_tensor ( x ) empty_tensor = tf . constant ( [ ] , dtype = x . dtype ) b = tf . convert_to_tensor ( b ) if b is not None else empty_tensor act_spec = activation_funcs [ act ] assert b . shape . rank == 1 and ( b . shape [ 0 ] == 0 or b . shape [ 0 ] == x . shape [ axis ] ) assert b . shape [ 0 ] == 0 or 0 <= axis < x . shape . rank if alpha is None : alpha = act_spec . def_alpha if gain is None : gain = act_spec . def_gain # Special cases. if act == 'linear' and b is None and gain == 1.0 : return x if act_spec . cuda_idx is None : return _fused_bias_act_ref ( x = x , b = b , axis = axis , act = act , alpha = alpha , gain = gain ) # CUDA kernel. cuda_kernel = _get_plugin ( ) . fused_bias_act cuda_kwargs = dict ( axis = axis , act = act_spec . cuda_idx , alpha = alpha , gain = gain ) # Forward pass: y = func(x, b). def func_y ( x , b ) : y = cuda_kernel ( x = x , b = b , ref = empty_tensor , grad = 0 , * * cuda_kwargs ) y . set_shape ( x . shape ) return y # Backward pass: dx, db = grad(dy, x, y) def grad_dx ( dy , x , y ) : ref = { 'x' : x , 'y' : y } [ act_spec . ref ] dx = cuda_kernel ( x = dy , b = empty_tensor , ref = ref , grad = 1 , * * cuda_kwargs ) dx . set_shape ( x . shape ) return dx def grad_db ( dx ) : if b . shape [ 0 ] == 0 : return empty_tensor db = dx if axis < x . shape . rank - 1 : db = tf . reduce_sum ( db , list ( range ( axis + 1 , x . shape . rank ) ) ) if axis > 0 : db = tf . reduce_sum ( db , list ( range ( axis ) ) ) db . set_shape ( b . shape ) return db # Second order gradients: d_dy, d_x = grad2(d_dx, d_db, x, y) def grad2_d_dy ( d_dx , d_db , x , y ) : ref = { 'x' : x , 'y' : y } [ act_spec . ref ] d_dy = cuda_kernel ( x = d_dx , b = d_db , ref = ref , grad = 1 , * * cuda_kwargs ) d_dy . set_shape ( x . shape ) return d_dy def grad2_d_x ( d_dx , d_db , x , y ) : ref = { 'x' : x , 'y' : y } [ act_spec . ref ] d_x = cuda_kernel ( x = d_dx , b = d_db , ref = ref , grad = 2 , * * cuda_kwargs ) d_x . set_shape ( x . shape ) return d_x # Fast version for piecewise-linear activation funcs. @ tf . custom_gradient def func_zero_2nd_grad ( x , b ) : y = func_y ( x , b ) @ tf . custom_gradient def grad ( dy ) : dx = grad_dx ( dy , x , y ) db = grad_db ( dx ) def grad2 ( d_dx , d_db ) : d_dy = grad2_d_dy ( d_dx , d_db , x , y ) return d_dy return ( dx , db ) , grad2 return y , grad # Slow version for general activation funcs. @ tf . custom_gradient def func_nonzero_2nd_grad ( x , b ) : y = func_y ( x , b ) def grad_wrap ( dy ) : @ tf . custom_gradient def grad_impl ( dy , x ) : dx = grad_dx ( dy , x , y ) db = grad_db ( dx ) def grad2 ( d_dx , d_db ) : d_dy = grad2_d_dy ( d_dx , d_db , x , y ) d_x = grad2_d_x ( d_dx , d_db , x , y ) return d_dy , d_x return ( dx , db ) , grad2 return grad_impl ( dy , x ) return y , grad_wrap # Which version to use? if act_spec . zero_2nd_grad : return func_zero_2nd_grad ( x , b ) return func_nonzero_2nd_grad ( x , b )