softmax_cross_entropy.py 文件源码

def softmax_log(x, use_cudnn):
    xp = cuda.get_array_module(x)
    if (xp != numpy and cuda.cudnn_enabled and use_cudnn and
            _cudnn_version >= 3000):
        oz_dtype = 'd' if x.dtype == 'd' else 'f'
        one = numpy.array(1, dtype=oz_dtype).ctypes
        zero = numpy.array(0, dtype=oz_dtype).ctypes
        handle = cudnn.get_handle()
        x_cube = x.reshape(x.shape[:2] + (-1, 1))
        desc = cudnn.create_tensor_descriptor(x_cube)
        y = xp.empty_like(x)
        libcudnn.softmaxForward(
            handle, _algorithm, _mode, one.data, desc.value,
            x_cube.data.ptr, zero.data, desc.value,
            y.data.ptr)
        return y
    else:
        log_z = logsumexp(xp, x)
        return x - log_z