nvidia.dali.fn.numba_function#

nvidia.dali.fn.numba_function(__input_0, __input_1=None, __input_2=None, __input_3=None, __input_4=None, __input_5=None, /, *, batch_processing=False, blocks=None, bytes_per_sample_hint=[0], in_types, ins_ndim, out_types, outs_ndim, preserve=False, run_fn, setup_fn=None, threads_per_block=None, device=None, name=None)#

调用 njit 编译的 Numba 函数。

run 函数应该是一个 Python 函数，可以在 Numba nopython 模式下编译。 接收单个输入并产生单个输出的函数应遵循以下定义

def run_fn(out0, in0)

其中 out0 和 in0 是输入和输出张量的 numpy 数组视图。 如果操作符配置为以批处理模式运行，则数组的第一个维度是样本索引。

请注意，该函数最多可以接受 6 个输入和 6 个输出。

此外，可选的 setup 函数计算输出的形状，以便 DALI 可以为输出分配内存，其定义如下

def setup_fn(outs, ins)

setup 函数为整个批次调用一次。 outs、ins 的第一个维度分别是输出/输入的数量。 第二个维度是样本索引。 例如，可以通过 outs[1][0] 访问第二个输出上的第一个样本。

如果未提供 setup 函数，则输出形状和数据类型将与输入相同。

注意

此操作符是实验性的，其 API 可能会在不另行通知的情况下更改。

警告

当 pipeline 启用条件执行时，必须采取额外的步骤来防止 run_fn 和 setup_fn 函数被 AutoGraph 重写。 有两种方法可以实现这一点

  1. 在全局范围内定义函数（即在 pipeline_def 范围之外）。

  2. 如果函数是另一个“工厂”函数的结果，则工厂函数必须在 pipeline 定义函数之外定义，并使用 @do_not_convert 修饰。

更多详细信息可以在 @do_not_convert 文档中找到。

示例 1

以下示例显示了一个简单的 setup 函数，它置换形状中维度的顺序。

def setup_change_out_shape(outs, ins):
    out0 = outs[0]
    in0 = ins[0]
    perm = [1, 0, 2]
    for sample_idx in range(len(out0)):
        for d in range(len(perm)):
            out0[sample_idx][d] = in0[sample_idx][perm[d]]

由于 setup 函数是为整个批次运行的，因此我们需要迭代并单独置换每个样本的形状。 对于 shapes = [(10, 20, 30), (20, 10, 30)]，它将生成 shapes = [(20, 10, 30), (10, 20, 30)] 的输出。

同时提供 run 函数

def run_fn(out0, in0):
    for i in range(in0.shape[0]):
        for j in range(in0.shape[1]):
            out0[j, i] = in0[i, j]

run 函数可以按样本或按批次工作，具体取决于 batch_processing 参数。

按批次工作的 run 函数可能如下所示

def run_fn(out0_samples, in0_samples):
    for out0, in0 in zip(out0_samples, in0_samples):
        for i in range(in0.shape[0]):
            for j in range(in0.shape[1]):
                out0[j, i] = in0[i, j]

按样本工作的 run 函数可能如下所示

def run_fn(out0, in0):
    for i in range(in0.shape[0]):
        for j in range(in0.shape[1]):
            out0[j, i] = in0[i, j]

此操作符允许序列输入并支持体积数据。

此操作符将不会从图中优化掉。

支持的后端

  • ‘cpu’

  • ‘gpu’

参数:

  • __input_0 (TensorList) – 算子的输入。

  • __input_[1..5] (TensorList, 可选) – 此函数最多接受 5 个可选的位置输入

关键字参数:

  • batch_processing (bool, 可选, 默认值 = False) –

    确定函数是为每个批次调用一次，还是为批次中的每个样本单独调用。

    当 batch_processing 设置为 True 时，该函数处理整个批次。 如果函数必须执行跨样本操作，并且如果可以重用大部分工作，则这是必要的。 对于其他用例，指定 False 并使用每样本处理函数允许操作符并行处理样本。

  • blocks (int 或 list of int, 可选) –

    3 项列表，指定用于每个网格的块数，用于

    执行 CUDA 内核

  • bytes_per_sample_hint (int 或 list of int, 可选, 默认值 = [0]) –

    输出大小提示，以每样本字节为单位。

    如果指定，则将预先分配驻留在 GPU 或页锁定主机内存中的操作符输出，以容纳此大小的样本批次。

  • in_types (DALIDataType 或 list of DALIDataType) – 输入的类型。

  • ins_ndim (int 或 list of int) – 输入形状应具有的维度数。

  • out_types (DALIDataType 或 list of DALIDataType) – 输出的类型。

  • outs_ndim (int 或 list of int) – 输出形状应具有的维度数。

  • preserve (bool, 可选, 默认值 = False) – 即使操作符的输出未使用，也阻止将其从图中删除。

  • run_fn (object) – 要调用的函数。 此函数必须在 Numba nopython 模式下工作。

  • setup_fn (object, 可选) – 设置输出形状的 Setup 函数。 此函数每个批次调用一次。 此外，此函数必须在 Numba nopython 模式下工作。

  • threads_per_block (int 或 list of int, 可选) –

    3 项列表，指定每项的线程数

    块用于执行 CUDA 内核

本页内容