compute.py 文件源码

def compute_convolution_nd(data, kernel, dimension: int, mode=ConvolutionMode.valid, element_wise: bool=False):
    mode_string = __get_convolution_mode_string(mode)
    result = []
    data_prefix_shape = data.shape[:-dimension]
    kernel_prefix_shape = kernel.shape[:-dimension]
    if element_wise:
        final_shape = element_wise_shape(data_prefix_shape, kernel_prefix_shape)[0]
        data = numpy.broadcast_to(data, final_shape + data.shape[-2:])
        kernel = numpy.broadcast_to(kernel, final_shape + kernel.shape[-2:])
        if final_shape:
            for index in array_index_traversal(final_shape):
                result.append(__compute_convolution_nd(data[index], kernel[index], dimension, mode_string))
            return numpy.array(result).reshape(final_shape + result[0].shape)
        else:
            return __compute_convolution_nd(data, kernel, dimension, mode_string)
    else:
        if kernel_prefix_shape:
            final_shape = data_prefix_shape + kernel_prefix_shape + basic_convolution_shape(data.shape[-dimension:], kernel.shape[-dimension:], dimension, mode_string)
            result = numpy.zeros(final_shape)
            for kernel_index in array_index_traversal(kernel_prefix_shape):
                sub_result_index = tuple(slice(None) for _ in data_prefix_shape) + kernel_index + tuple(slice(None) for _ in range(dimension))
                result[sub_result_index] = __compute_convolution_nd(data, kernel[kernel_index], dimension, mode_string)
            return result
        else:
            return __compute_convolution_nd(data, kernel, dimension, mode_string)