跳转至主要内容
Version: v1.1.0

与外部数组进行交互

尽管 Taichi field 主要被用于 Taichi scope,在一些情况下在 Python scope 中高效地操作场数据是非常有帮助的。

我们提供了不同的接口用以在 Taichi field 和外部数组之间拷贝数据。 外部数组是指NumPy 数组、PyTorch tensor 或 Paddle tensor。 我们来看一下最常见的用法:与 NumPy 数组的交互。

使用 to_numpy()Taichi 场中的数据导出至 NumPy 数组。 这允许我们将计算结果导出至其他支持 NumPy 的 Python 包中,例如 matplotlib

@ti.kernel
def my_kernel():
for i in x:
x[i] = i * 2

x = ti.field(ti.f32, 4)
my_kernel()
x_np = x.to_numpy()
print(x_np) # np.array([0, 2, 4, 6])

使用 from_numpy()NumPy 数组导入至 Taichi 场。 这允许我们使用 NumPy 数组来初始化 Taichi 场。

x = ti.field(ti.f32, 4)
x_np = np.array([1, 7, 3, 5])
x.from_numpy(x_np)
print(x[0]) # 1
print(x[1]) # 7
print(x[2]) # 3
print(x[3]) # 5

同样,Taichi field 可以被 导入并导出到 PyTorch tensor

@ti.kernel
def my_kernel():
for i in x:
x[i] = i * 2

x = ti.field(ti.f32, 4)
my_kernel()
x_torch = x.to_torch()
print(x_torch) # torch.tensor([0, 2, 4, 6])

x.from_numpy(torch.tensor([1, 7, 3, 5]))
print(x[0]) # 1
print(x[1]) # 7
print(x[2]) # 3
print(x[3]) # 5

Taichi field 可以被 导入并导出到 Paddle tensor

@ti.kernel
def my_kernel():
for i in x:
x[i] = i * 2

x = ti.field(ti.f32, 4)
my_kernel()
x_paddle = x.to_paddle()
print(x_paddle) # paddle.Tensor([0, 2, 4, 6])

x.from_numpy(paddle.to_tensor([1, 7, 3, 5]))
print(x[0]) # 1
print(x[1]) # 7
print(x[2]) # 3
print(x[3]) # 5

调用 to_torch() 时,请使用 device 参数指定 Taichi field 需要导出到的 PyTorch 设备:

x = ti.field(ti.f32, 4)
x.fill(3.0)
x_torch = x.to_torch(device="cuda:0")
print(x_torch.device) # device(type='cuda', index=0)

对于 Paddle, 请使用 paddle.CPUPlace()paddle.CUDAPlace(n) 制定设备。其中 n 是可选的 id, 默认值为 0。

外部数组形状

Taichi field 的尺寸及相应的 NumPy 数组、PyTorch tensor 或是 Paddle Tensor 的尺寸由以下规则相关联:

  • 对于标量 field, NumPy 数组、PyTorch tensor 或 Paddle Tensor 与Taichi field 的形状相同
field = ti.field(ti.i32, shape=(256, 512))
field.shape # (256, 512)

array = field.to_numpy()
array.shape # (256, 512)

field.from_numpy(array) # the input array must be of shape (256, 512)
  • 对于向量 field,如果向量是 n维的, 那么 NumPy 数组、 PyTorch tensor 或 Paddle Tensor 的形状应为 (*field_shape, vector_n)
field = ti.Vector.field(3, ti.i32, shape=(256, 512))
field.shape # (256, 512)
field.n # 3

array = field.to_numpy()
array.shape # (256, 512, 3)

field.from_numpy(array) # the input array must be of shape (256, 512, 3)
  • 对于矩阵 field,如果矩阵是 n乘以m (n x m), 那么 NumPy 数组、 PyTorch tensor 或 Paddle Tensor 的形状应为 (formfield_shape, matrix_n, matrix_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) # the input array must be of shape (256, 512, 3, 4)
  • 对于结构体 field,外部数组将被导出为一个 NumPy 数组或 PyTorch tensor 或 Paddle Tensor 的字典,这个字典的键为结构体 field 中的结构体成员名,键对应的值为结构体 field 中的结构体成员数组。 嵌套结构将导出为嵌套字典:
field = ti.Struct.field({'a': ti.i32, 'b': ti.types.vector(float, 3)} 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) # the input array must have the same keys as the field

使用外部数组作为 Taichi 内核的参数

使用类型提示 ti.types.ndarray() 将外部数组作为 kernel 的参数传入。 For example:

import taichi as ti
import numpy as np

ti.init()

n, m = 4, 7
a = np.empty(shape=(n, m), dtype=np.int32)


@ti.kernel
def test_numpy(arr: ti.types.ndarray()):
# You can access the shape of the passed array in the kernel
for i in range(arr.shape[0]):
for j in range(arr.shape[1]):
arr[i, j] += i + j


for i in range(n):
for j in range(m):
a[i, j] = i * j

test_numpy(a)

for i in range(n):
for j in range(m):
assert a[i, j] == i * j + i + j

请注意,外部数组中的元素必须用一个方括号索引。 这与Taichi 的向量或矩阵 field 不同,Taichi 的 field 和矩阵元素可以分别索引:

@ti.kernel
def copy_vector(x: ti.template(), y: ti.types.ndarray()):
for i, j in ti.ndrange(n, m):
for k in ti.static(range(3)):
y[i, j, k] = x[i, j][k] # correct
# y[i][j][k] = x[i, j][k] incorrect
# y[i, j][k] = x[i, j][k] incorrect

另外,Taichi kernel 中的外部阵列使用它们自己的 物理内存布局 索引。 对于PyTorch用户来说,这意味着在传入Taichi kernel 之前,PyTorch tensor 需要保持连续

@ti.kernel
def copy_scalar(x: ti.template(), y: ti.types.ndarray()):
for i, j in x:
y[i, j] = x[i, j]

x = ti.field(dtype=int, shape=(3, 3))
y = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
y = y.T # Transposing the tensor returns a view of the tensor which is not contiguous
copy(x, y) # error!
copy(x, y.clone()) # correct
copy(x, y.contiguous()) # correct