Taichi 外部数组数据交互实战：NumPy / PyTorch / Paddle 与 Taichi Field、Kernel 参数的完整互通指南

Taichi 外部数组数据交互实战NumPy / PyTorch / Paddle 与 Taichi Field、Kernel 参数的完整互通指南【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi本文是 Taichi 语言中“外部数组External Arrays数据交互”的完整技术指南内容以仓库文档 docs/lang/articles/basic/external.md 为骨架结合 python/taichi/lang/field.py、python/taichi/types/ndarray_type.py、python/taichi/_kernels.py 等源码实现与 tests/python/test_field.py 测试用例展开。读完本文你将掌握三件事如何把 NumPy/PyTorch/Paddle 数组按值复制进 Taichi field 或按引用传递给 kernel标量/向量/矩阵/结构体 field 与外部数组之间严格的形状匹配规则以及外部数组作为 kernel 参数时的索引方式、内存布局约束与常见 FAQ。概述Taichi 与外部数组的两种交互方式Taichi 目前支持三种外部数组NumPy 数组numpy.ndarray、PyTorch 张量torch.Tensor和 Paddle 张量paddle.Tensor。由于 NumPy 数组是 Taichi 中最常用的外部数组本文以其为例讲解其余两种的操作步骤完全一致。把一个 NumPy 数组arr导入 Taichi scope有两种途径二者语义截然不同创建 shape 与 dtype 匹配的 Taichi field调用f.from_numpy(arr)将数据复制进 field。该方式为按值拷贝Taichi 会维护一份自己的数据副本此后外部数组与原 field 互不影响。此方式适合“原始数组在 Taichi scope 内被频繁访问”的场景例如纹理采样时反复读取同一份数据。用ti.types.ndarray()作为类型注解将arr作为参数传入 kernel 或 Taichi 函数。该方式为按引用传递参数不会产生副本kernel 内部对参数的任何修改都会直接反映到原始数组arr上。此方式适合“kernel 需要直接处理原始数组”的场景例如存储结果、做滤波等可以避免一次额外的数据拷贝开销。import taichi as ti import numpy as np ti.init() a np.zeros((5, 5)) ti.kernel def test(a: ti.types.ndarray()): for i in range(a.shape[0]): # 并行 for 循环 for j in range(a.shape[1]): a[i, j] i j test(a) print(a) # 原数组 a 已被直接修改注意from_numpy()/from_torch()可以接收任意NumPy 数组或 torch Tensor——无论是否连续contiguousTaichi 都会自行管理自己的数据副本。然而当数组作为参数传入 kernel时只支持连续contiguous的NumPy 数组或 torch Tensor详见后文“物理内存布局”一节。NumPy 数组与 Taichi field 的数据互转从 NumPy 导入 fieldfrom_numpy()导入前必须保证 field 与数组的 shape 完全一致import taichi as ti import numpy as np x ti.field(float, shape(3, 3)) a np.arange(9).reshape(3, 3).astype(np.int32) x.from_numpy(a) print(x) #[[0 1 2] # [3 4 5] # [6 7 8]]上面示例中标量 fieldx与数组a的 shape 同为(3, 3)操作成功。如果 shape 不匹配该操作会直接抛错——这一点在源码中有明确的硬校验ScalarField._from_external_arr()会先比较维数、再逐维比较大小不匹配即抛出ValueError见 python/taichi/lang/field.pyif len(self.shape) ! len(arr.shape): raise ValueError(fti.field shape {self.shape} does not match f the numpy array shape {arr.shape}) for i, _ in enumerate(self.shape): if self.shape[i] ! arr.shape[i]: raise ValueError(...)仓库测试 tests/python/test_field.py 中的test_scalr_field_from_numpy_with_mismatch_shape专门验证了该行为向 shape 不匹配的 field 调用from_numpy时测试断言抛出ValueError。此外field 的 dtype 与数组的 dtype 最好一致若不一致会发生隐式类型转换详见 类型系统文档。另外from_numpy()内部会检查arr.flags.c_contiguous对非 C 连续数组先执行np.ascontiguousarray()再拷贝见 python/taichi/lang/field.py因此对非连续数组也能安全工作。从 field 导出到 NumPyto_numpy()将 field 数据导出为 NumPy 数组调用to_numpy()即可arr x.to_numpy() #array([[0, 1, 2], # [3, 4, 5], # [6, 7, 8]], dtypeint32)从源码看python/taichi/lang/field.pyto_numpy()会先根据 field dtype 经to_numpy_type()映射出对应的 NumPy dtype映射表见 python/taichi/lang/util.py分配零数组后调用内部 kerneltensor_to_ext_arr定义于 python/taichi/_kernels.py逐元素拷贝最后执行runtime_ops.sync()同步后返回。to_numpy()还支持可选的dtype参数用于指定返回数组的目标 dtype。PyTorch / Paddle 张量与 Taichi field 的数据互转PyTorchfrom_torch()与to_torch(device...)PyTorch 张量与 field 的互转和 NumPy 完全类似from_torch()负责导入to_torch()负责导出。唯一区别是to_torch()多一个必填参数device用于指定 PyTorch 设备tensor x.to_torch(devicecuda:0) print(tensor.device) # device(typecuda, index0)源码中to_torch()在指定 device 上分配torch.zeros后调用tensor_to_ext_arr拷贝见 python/taichi/lang/field.py。同时注意from_torch()在基类 python/taichi/lang/field.py 中会先调用arr.contiguous()再导入因此非连续张量也可安全导入 field。Paddlefrom_paddle()与to_paddle(place...)Paddle 侧的设备指定方式与 PyTorch 不同需要显式传入一个place对象paddle.CPUPlace()或paddle.CUDAPlace(n)其中n是可选设备 ID默认值为 0import paddle device paddle.CPUPlace() tensor x.to_paddle(devicedevice)从源码看python/taichi/lang/field.pyto_paddle()先用to_paddle_type()见 python/taichi/lang/util.py完成 dtype 映射再经paddle.to_tensor(..., placeplace)放置到指定设备。外部数组形状匹配规则在ti.field、ti.Vector.field、ti.Matrix.field与外部数组之间传数时必须保证两侧形状对齐。规则按 field 类型分四种情况统一归纳如下规则 1标量 field —— 形状完全相等导入/导出标量 field 时外部数组的 shape必须等于标量 field 的 shapefield ti.field(int, shape(256, 512)) field.shape # (256, 512) array field.to_numpy() array.shape # (256, 512) field.from_numpy(array) # 输入数组 shape 必须为 (256, 512)对应关系示意field.shape[1]array.shape[1] (512) ┌───────────────────────┐ ┌ ┌───┬───┬───┬───┬───┬───┐ ┐ │ │ │ │ │ │ │ │ │ │ ├───┼───┼───┼───┼───┼───┤ │ field.shape[0]array.shape[0] │ │ │ │ │ │ │ │ │ (256) │ ├───┼───┼───┼───┼───┼───┤ │ │ │ │ │ │ │ │ │ │ └ └───┴───┴───┴───┴───┴───┘ ┘规则 2n 维向量 field —— 形状为(*field_shape, n)导入/导出 n 维向量 field 时外部数组的 shape 为(*field_shape, n)即把向量维度拼在 field 形状之后作为最后一维field ti.Vector.field(3, int, shape(256, 512)) field.shape # (256, 512) field.n # 3 array field.to_numpy() array.shape # (256, 512, 3) field.from_numpy(array) # 输入数组 shape 必须为 (256, 512, 3)对应关系示意field.shape[1]array.shape[1] (512) ┌─────────────────────────────┐ ┌ ┌─────────┬─────────┬─────────┐ ┐ │ │[*, *, *]│[*, *, *]│[*, *, *]│ │ │ ├─────────┼─────────┼─────────┤ │ field.shape[0]array.shape[0] │ │[*, *, *]│[*, *, *]│[*, *, *]│ │ [*, *, *] (256) │ ├─────────┼─────────┼─────────┤ │ └───────┘ │ │[*, *, *]│[*, *, *]│[*, *, *]│ │ narray.shape[2]3 └ └─────────┴─────────┴─────────┘ ┘规则 3n×m 矩阵 field —— 形状为(*field_shape, n, m)导入/导出 n 行 m 列矩阵 field 时外部数组的 shape 为(*field_shape, n, m)矩阵的两个维度依次拼在最后field ti.Matrix.field(3, 4, ti.i32, shape(256, 512)) field.shape # (256, 512) field.n # 3 field.m # 4 array field.to_numpy() array.shape # (256, 512, 3, 4) field.from_numpy(array) # 输入数组 shape 必须为 (256, 512, 3, 4)顺带一提MatrixField.to_numpy()还提供keep_dims与dtype两个可选参数keep_dimsTrue时即使矩阵维度为 1×1、1×n、n×1 也始终保留 n2 维输出keep_dimsFalse默认则会跳过大小为 1 的矩阵维度见 python/taichi/lang/matrix.py。规则 4结构体 field —— 外部数组以“字典”形式出现导入/导出 struct field 时外部数组表现为字典键key为结构体成员名值value为对应成员数组。嵌套的结构体导出为嵌套字典field ti.Struct.field({a: ti.i32, b: ti.types.vector(3, float)}, shape(256, 512)) field.shape # (256, 512) array_dict field.to_numpy() array_dict.keys() # dict_keys([a, b]) array_dict[a].shape # (256, 512) array_dict[b].shape # (256, 512, 3) field.from_numpy(array_dict) # 输入字典的键必须与 field 成员一一对应从实现看StructField.to_numpy()本质是对每个成员分别调用v.to_numpy()再组装成字典from_numpy()则遍历self._items逐成员调用对应的from_numpy见 python/taichi/lang/struct.py。同理struct field 也支持from_torch/to_torch/from_paddle/to_paddle的字典形式。外部数组作为 Taichi kernel 参数入门示例ti.types.ndarray()上文已经展示了最基础的用法。这里的核心是类型注解ti.types.ndarray()。它在源码中的实现为 python/taichi/types/ndarray_type.py 的NdarrayType类模块内同时导出别名ndarray NdarrayTypeti.kernel def test(a: ti.types.ndarray()): for i in range(a.shape[0]): # 并行 for 循环 for j in range(a.shape[1]): a[i, j] i jNdarrayType支持若干可选参数以约束传入数组参数含义dtype元素类型PrimitiveType/VectorType/MatrixType不指定时不做校验ndim数组维度数当前对外部数组暂忽略element_dim元素维度0 表示标量、1 表示向量、2 表示矩阵element_shape每个元素的形状向量为一维元组、矩阵为二维元组needs_grad是否要求传入带梯度的数组boundary越界访问边界处理默认unsafe当注解中指定了 dtype/ndim/needs_grad 时kernel 编译期会调用NdarrayType.check_matched()见 python/taichi/types/ndarray_type.py对实际传入的数组做逐一校验元素类型不匹配抛TypeError维度数不匹配抛ValueError并输出含参数名与期望值的明确报错信息。进阶示例离散拉普拉斯算子假设a、b是两个 shape 与 dtype 相同的二维数组。我们要对a中每个格子(i, j)计算其值与上下左右四个邻居均值的差并把结果写入b的对应位置。为简化边界格子邻居不足四个被排除在外。这个操作通常称为离散拉普拉斯算子Discrete Laplace Operatorb[i, j] a[i, j] - (a[i-1, j] a[i, j-1] a[i1, j] a[i, j1]) / 4这类操作即使借助 NumPy 的向量化写法也往往很慢切片拼接、内存搬运开销大b[1:-1, 1:-1] ( a[ :-2, 1:-1] a[1:-1, :-2] a[1:-1, 2:] a[2: , 1:-1])而 Taichi 只需一个并行for循环即可完成同样的事情ti.kernel def test(a: ti.types.ndarray(), b: ti.types.ndarray()): # 假设 a、b shape 相同 H, W a.shape[0], a.shape[1] for i, j in ti.ndrange(H, W): # 一个并行 for 循环 if 0 i H - 1 and 0 j W - 1: b[i, j] a[i, j] - (a[i-1, j] a[i, j-1] a[i1, j] a[i, j1]) / 4这段代码不仅比 NumPy 版本更易读而且在 CPU 后端上运行速度也明显更快并行循环自动调度无需逐格手动切片。外部数组的索引方式与物理内存布局外部数组元素必须用“单一对方括号”索引这与 Taichi 向量/矩阵 field 的“成员与元素分开索引”截然不同x ti.Vector.field(3, float, shape(5, 5)) y np.random.random((5, 5, 3)) ti.kernel def copy_vector(x: ti.template(), y: ti.types.ndarray()): for i, j in ti.ndrange(5, 5): for k in ti.static(range(3)): y[i, j, k] x[i, j][k] # 正确 # y[i][j][k] x[i, j][k] # 错误 # y[i, j][k] x[i, j][k] # 错误另一个关键约束是kernel 内部是按外部数组的“物理内存布局”索引的。对 PyTorch 用户来说这意味着张量在传入 kernel 前必须确保是连续的contiguous。典型反例是转置操作y.T返回的是原张量的视图view并不连续直接传入会出错x ti.field(dtypeint, shape(3, 3)) y torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) y y.T # 转置后得到的是不连续的 view ti.kernel def copy_scalar(x: ti.template(), y: ti.types.ndarray()): for i, j in x: y[i, j] x[i, j] # copy(x, y) # 错误y 不连续 copy(x, y.clone()) # 正确 copy(x, y.contiguous()) # 正确与此形成对照的是from_numpy()/from_torch()它们内部会主动处理非连续数组np.ascontiguousarray()/arr.contiguous()所以导入 field 时无需用户手动转连续只有“作为 kernel 参数按引用传递”这一路径才要求数组必须连续。FAQ可以用ti.kernel加速任意 NumPy 函数吗不能。与其他 Python 加速框架如 Numba不同Taichi不编译 NumPy 函数在 Taichi scope 内调用 NumPy 函数是不支持的import numpy as np ti.kernel def invalid_sum(arr: ti.types.ndarray()): total np.sum(arr) # 不支持 ...如果确实需要用到 Taichi 中没有对应实现的 NumPy 函数正确的做法是在 Python scope 中照常调用该函数再把处理好的数组通过ti.types.ndarray()传入 Taichi kernelarr np.random.random(233) indices np.argsort(arr) # arr 是 NumPy ndarray在 Python scope 中调用 ti.kernel def valid_example(arr: ti.types.ndarray(), indices: ti.types.ndarray()): min_element arr[indices[0]] ...这种“Python scope 预处理 kernel 内并行计算”的组合正是 Taichi 与外部生态协作的推荐模式。小结外部数组交互是 Taichi 日常开发中使用频率最高的能力之一核心要点可概括为四条按值复制用from_numpy/from_torch/from_paddle按引用传参用ti.types.ndarray()二者语义不同、适用场景不同形状匹配有固定规则标量 field 完全相等、向量 field 尾部加n、矩阵 field 尾部加n, m、struct field 用键名一一对应的字典kernel 参数按物理内存布局索引单一对方括号取元素PyTorch 张量必须先contiguous()Taichi 不编译 NumPy 函数需要时在 Python scope 里先算好再交给 kernel 做并行处理。如需进一步阅读可继续查看 基础 field 文档、布局文档、ndarray 文档 与 类型系统文档以及仓库中 tests/python/test_field.py、tests/python/test_argument.py、tests/python/test_get_external_tensor_shape.py 等测试文件中的更多实操用例。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考