参考文档
1. 初始化 Taichi
ti.init(arch = ti.cuda)
推荐使用:
ti.init(arch = ti.gpu)
会自动识别 cuda
/ metal
/ opengl
2. 数据类型
# Signed integers:
ti.i8/i16/i32/i64
# Unsigned integers:
ti.u8/u16/u32/u64
# Float-point numbers:
ti.f32/f64
NOTE:
- 布尔类型暂时使用
ti.i32
- opengl 不支持64位的整型和浮点型,详情查看文档
3. 成员变量
3.1 Tensors(张量)
- Taichi 中的张量
Tensors
本质是一个多维数组 - 一个张量元素可以是:
- 零维张量
var
:标量形式 - 一维张量
ti.Vector
:向量形式 - 二维张量
ti.Matrix
:矩阵形式
- 零维张量
- 张量元素总是通过
a[i,j,k]
语法访问。 - Taichi 中的张量和矩阵是两个概念
- 张量可以是空间稀疏的
import taichi as ti
ti.init()
a = ti.var(dt=ti.f32, shape=(42, 63))
# A tensor of 42x63 scalars
# 标量的张量,其中 42x63 个元素
b = ti.Vector(3, dt=ti.f32, shape=4)
# A tensor of 4x 3D vectors
# 4个三维向量组成的张量
# 张量的元素是三维向量
# Vector(向量的维度, dt = 数据类型, shape = 张量的大小:向量的个数)
C = ti.Matrix(2, 2, dt=ti.f32, shape=(3, 5))
# A tensor of 3x5 2x2 matrices
# 张量的元素是 2*2 的矩阵
# 张量的形状是 3*5 ,其中的每个元素是 2*2 的矩阵
loss = ti.var(dt=ti.f32, shape=())
# A (0−D) tensor of a single scalar
# 零维张量:只有单个元素的张量,且该元素是个标量
a[3, 4] = 1
print('a[3,4] =', a[3, 4])
# ”a[3,4] = 1.000000”
b[2] = [6, 7, 8]
print(' b[0] =', b[0][0], b[0][1], b[0][2])
# print (b[0]) is not yet supported .
loss[None] = 3
print(loss[None])
# 3
# 在访问 loss 的时候可以将下标设置为 None
# 注意不要直接 loss = 3 ,这样变量 loss 的类型将不再是 ti.var
3.2 field
Note :在 taichi 0.6.22
往后用 field
进行了替换:
pip install --upgrade taichi==0.6.22
import taichi as ti
ti.init()
a = ti.field(dtype=ti.f32, shape=(42, 63))
# 标量的张量,其中 42x63 个元素
b = ti.Vector.field(3, dtype=ti.f32, shape=4)
# 4个三维向量组成的张量
# 张量的元素是三维向量
# Vector(向量的维度, dt = 数据类型, shape = 张量的大小:向量的个数)
C = ti.Matrix.field(2, 2, dtype=ti.f32, shape=(3, 5))
# 张量的元素是 2*2 的矩阵
# 张量的形状是 3*5 ,其中的每个元素是 2*2 的矩阵
loss = ti.field(dtype=ti.f32, shape=())
# 零维张量:只有单个元素的张量,且该元素是个标量
a[3, 4] = 1
print('a[3,4] =', a[3, 4])
# ”a[3,4] = 1.000000”
b[2] = [6, 7, 8]
print(' b[0] =', b[0][0], b[0][1], b[0][2])
# print (b[0]) is not yet supported .
loss[None] = 3
print(loss[None])
# 3
# 在访问 loss 的时候可以将下标设置为 None
# 注意不要直接 loss = 3 ,这样变量 loss 的类型将不再是 ti.var
带 field
是全局变量,不带的是局部变量:
x = ti.Vector.field(3, dtype=ti.f32, shape=(128, 512)) # global
@ti.kernel
def foo():
a = ti.Vector([0.2, 0.4]) # local
4. Kernels
-
Taichi 中的
Kernels
是用来计算的函数,用法类似于 python 中的Functions
-
Taichi 中的
Kernels
有编译过程,提高代码运行速度 -
静态类型,在GPU上并行运行,且可微分
-
使用
Kernels
必须装饰@ti.kernel
-
Kernel 的参数和返回值必须声明类型(Taichi 本身是强类型语言)
@ti.kernel
def hello(i: ti.i32):
a = 40
print('Hello world!', a + i)
hello(2) # "Hello world! 42"
@ti.kernel
def calc() -> ti.i32:
s = 0
for i in range(10):
s += i
return s # 45
5. Functions
- Taichi 中的
Functions
可以被 Taichi 中的Kernels
调用 - Taichi 中的
Kernels
可以被 Taichi 中的Functions
调用
- Taichi 中的
Functions
不能被 python 中的Functions
调用 - Taichi 中的
Kernels
可以被 python 中的Functions
调用
- Taichi 中的
Functions
可以被 Taichi 中的Functions
调用 - Taichi 中的
Kernels
不能被 Taichi 中的Kernels
调用
Note :
编写代码的过程中尽量遵循一个原则:
- 将
ti.func
作为所有数值处理的子块,其内联的特性也将提高其运算性能; - 将
ti.kernel
仅作为向外与 pyhton 生态交互的接口,其作用是分布ti.func
任务,同时作为程序的外壳暴露给诸如渲染循环等结构。
@ti.func
def triple(x):
return x * 3
@ti.kernel
def triple_array:
for i in range(128):
a[i] = triple(a[i])
NOTE:
- Taichi 的
Functions
将是强制 inline 的。目前,不允许递归。 - Taichi 的
Functions
最多只能包含一个return
语句。
6. 标准数据库
ti.sin(x)
ti.cos(x)
ti.asin(x)
ti.acos(x)
ti.atan2(y, x)
ti.cast(x, data_type)
ti.sqrt(x)
ti.floor(x)
ti.ceil(x)
ti.inv(x)
ti.tan(x)
ti.tanh(x)
ti.exp(x)
ti.log(x)
ti.random(data_type)
# 传入数据类型,返回数据类型中的一个数
# python 的符号
a + b
a / b # 结果为浮点数
a // b # 结果为整型
# python 自带的库
abs(x)
int(x)
float(x)
max(x, y, ...)
min(x, y, ...)
x ** y
# x^y
7. 矩阵与线性代数
ti.Matrix
只能用于小矩阵(例如 3×3)- 如果使用 64×64 数组,建议使用一个标量的 2D 张量:
ti.var(dt=ti.f32, shape=(64, 64))
ti.Vector
向量只能用于一列的元素- 区分 Hadamard 积
*
(也称对应元素乘积,A 和 B 对应元素相乘获得新的矩阵)和矩阵积@
(矩阵乘法)
以下方法是返回一个新的矩阵:
A.transpose()
# 矩阵的转置
A.inverse()
# 矩阵的逆
A.trace()
# 矩阵的迹
A.determinant(type)
# 矩阵的行列式(需要设置数据类型)
A.normalized()
# 矩阵标准化
A.cast(type)
# 重置矩阵中的数据类型
R, S = ti.polar_decompose(A, ti.f32)
# 矩阵的极分解
U, sigma , V = ti.svd(A, ti.f32)
# 矩阵的奇异值分解
# sigma 是对角矩阵
ti.sin(A)/cos(A)... (element -wise)
# 将矩阵中的每个元素进行 sin 计算
# 所有的标量运算符均可以在矩阵上运算
8. 并行 for 循环
因为 Taichi 中的循环有两种形式
-
Range-for loops
:- 类似于 python 中的 for 循环;
- 只是在最外层的范围内使用它时会被并行化;
- 循环内并行执行的代码块顺序是不确定的;
- 循环的范围可以嵌套。
-
Struct-for loops
:迭代(稀疏)张量元素。
8.1 Range-for loops
在 Taichi 的 kernel
中最外层作用域的 For 循环被自动并行化。
@ti.kernel
def fill():
for i in range(10): # 并行执行,0 <= i < 4
x[i] += i
s = 0
for j in range(5): # 在最外层for循环的每个线程中串行执行
s += j
y[i] = s
@ti.kernel
def fill_3d():
# Parallelized for all 3 <= i < 8, 1 <= j < 6, 0 <= k < 9
# ti.ndrange 可以定义上下限
for i, j, k in ti.ndrange((3, 8), (1, 6), 9):
x[i, j, k] = i + j + k
NOTE :只有最外层范围的循环可以并行化;但循环若不是在最外层(例如最外层还有 if 语句),则无法自动并行。
@ti.kernel
def foo():
for i in range(10): # Parallelized :-)
...
@ti.kernel
def bar(k: ti.i32):
if k > 42:
for i in range(10): # Serial :-(
...
8.2 Struct-for loops
import taichi as ti
ti.init(arch=ti.gpu)
n = 320
pixels = ti.var(dt=ti.f32, shape=(n * 2, n))
@ti.kernel
def paint(t: ti.f32):
for i, j in pixels: # Parallized over all pixels
pixels[i, j] = i * 0.001 + j * 0.002 + t
paint(0.3)
Struct-for loops 遍历所有张量坐标,即:
(0, 0), (0, 1), (0, 2), …, (0, 319)
(1, 0), …
…
…, (639, 319).
9. 原子操作
- 在 Taichi 中,自赋值(例如
x[i]+=1
)是自动原子化的 - 在并行的修改全局变量时,请确保使用原子操作。(例如,统计x中的所有元素)
@ti.kernel
def sum():
for i in x:
# Approach 1: OK
total[None] += x[i]
# Approach 2: OK
ti.atomic_add(total[None], x[i])
# Approach 3: Wrong result (the operation is not atomic.)
# 存在非原子的 读加写 的操作
total[None] = total[None] + x[i]
10. Taichi-scope v.s. python-scope
Taichi-scope :任何被 ti.kernel
和 ti.func
修饰的函数的函数体。
python-scope :任何不在 Taichi 里面的函数体。
区别:
在 Taichi-scope 里的代码会被 Taichi 编译,并且在并行设备上运行。
在 python-scope 的代码不会被编译,仅仅只会被 python 解释器运行。
import taichi as ti
ti.init ()
a = ti . var (dt=ti. f32, shape =(42, 63)) # A tensor of 42x63 scalars
b = ti . Vector (3 , dt=ti . f32 , shape=4) # A tensor of 4x 3D vectors
C = ti . Matrix (2 , 2 , dt=ti.f32 , shape =(3, 5) ) # A tensor of 3x5 2x2 matrices
@ti.kernel
def foo () :
a[3,4] = 1
print ('a[3,4] =',a[3,4])
# ”a[3,4] = 1.000000”
b[2] = [6,7,8]
print('b[0]=',b[0],',b[2] =',b[2])
# ”b[0] = [[0.000000],[0.000000],[0.000000]],b[2]=[[6.000000],[7.000000],[8.000000]]”
C[2,1][0,1] = 1
print('C[2,1] =',C[2,1])
# C[ 2 , 1] = [[0.000000 , 1.000000] , [0.000000 , 0.000000]]
foo()
11. 一个 Taichi 程序的几个阶段
- 初始化:
ti.init(...)
- Tensor 定义(到计算阶段就无法继续定义):
ti.var
,ti.Vector
,ti.Matrix
- 计算阶段 (启动 kernels)
- 可选: 重新启动太极系统(清除内存,销毁所有变量和 kernels ):
ti.reset()
import taichi as ti
ti.init(arch=ti.gpu)
n = 320
pixels = ti.var(dt=ti.f32, shape=(n * 2, n))
@ti.func
def complex_sqr (z) :
return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2 ])
@ti.kernel
def paint (t:ti.f32):
for i,j in pixels: # Parallized over all pixels
c = ti.Vector([-0.8, ti.cos(t) * 0.2])
z = ti.Vector([ i/n - 1 , j/n - 0.5 ]) * 2
iterations = 0
while z.norm() < 20 and iterations < 50:
z = complex_sqr(z) + c
iterations += 1
pixels[i,j] = 1 - iterations * 0.02
gui = ti.GUI("Julia Set" , res = (n * 2, n))
for i in range(1000000) :
paint(i * 0.03)
gui.set_image(pixels)
gui.show()
12. Debug 模式
在调试模式下初始化Taichi,启用绑定检查程序(仅限CPU):ti.init(debug=True, arch=ti.cpu)
import taichi as ti
ti.init(debug=True , arch=ti.cpu)
a = ti.var(ti.i32, shape=(10))
b = ti.var(ti.i32, shape=(10))
@ti.kernel
def shift():
for i in range(10):
a[i] = b[i + 1]
# 会出现越界的情况,仅仅在debug 的情况下检测出越界
# Runtime error in debug mode
shift()
12.1 关闭优化
关闭优化,让编译速度变得更快
ti.core.toggle_advanced_optimization(False)
13. ODOP
在面向数据编程中加入面向对象特性
Demo :
ti example odop_solar
a = G M r / ∥ r ∥ 2 3 \bold{a} = GM\bold{r} / \Vert \bold{r} \Vert_2^3 a=GMr/∥r∥23
import taichi as ti
import math
ti.init()
@ti.data_oriented
class SolarSystem:
def __init__(self, n, dt): # Initializer of the solar system simulator
self.n = n
self.dt = dt
self.x = ti.Vector.field(2, dtype=ti.f32, shape=n)
self.v = ti.Vector.field(2, dtype=ti.f32, shape=n)
self.center = ti.Vector.field(2, dtype=ti.f32, shape=())
@staticmethod
@ti.func
def random_vector(radius): # Create a random vector in circle
theta = ti.random() * 2 * math.pi
r = ti.random() * radius
return r * ti.Vector([ti.cos(theta), ti.sin(theta)])
@ti.kernel
def initialize_particles(self):
# (Re)initialize particle position/velocities
for i in range(self.n):
offset = self.random_vector(0.5)
self.x[i] = self.center[None] + offset # Offset from center
self.v[i] = [-offset.y, offset.x] # Perpendicular to offset
self.v[i] += self.random_vector(0.02) # Random velocity noise
self.v[i] *= 1 / offset.norm()**1.5 # Kepler's third law
@ti.func
def gravity(self, pos): # Compute gravity at pos
offset = -(pos - self.center[None])
return offset / offset.norm()**3
@ti.kernel
def integrate(self): # Semi-implicit Euler time integration
for i in range(self.n):
self.v[i] += self.dt * self.gravity(self.x[i])
self.x[i] += self.dt * self.v[i]
def render(self, gui): # Render the scene on GUI
gui.circle([0.5, 0.5], radius=10, color=0xffaa88)
gui.circles(solar.x.to_numpy(), radius=3, color=0xffffff)
solar = SolarSystem(8, 0.0001)
solar.center[None] = [0.5, 0.5]
solar.initialize_particles()
gui = ti.GUI("Solar System", background_color=0x0071a)
while gui.running:
if gui.get_event() and gui.is_pressed(gui.SPACE):
solar.initialize_particles() # reinitialize when space bar pressed.
for i in range(10): # Time integration
solar.integrate()
solar.render(gui)
gui.show()
14. Templates
14.1 template 元编程
@ti.kernel
def offset(x: ti.template(), y: ti.template(), c: ti.f32):
for i in x:
y[i] = x[i] + c
当使用 ti.template()
时,每次调用该 kernel 都会实例化一个对应的 Kernel
import taichi as ti
ti.init()
@ti.kernel
def hello(i: ti.template()):
print(i)
for i in range(100):
hello(i) # 100 different kernels will be created
@ti.kernel
def world(i: ti.i32):
print(i)
for i in range(100):
world(i) # The only instance will be reused现
复制不同阶的 Tensor :
@ti.kernel
def copy(x: ti.template(), y: ti.template()):
for I in ti.grouped(y):
x[I] = y[I]n
NOTE :
- 其中使用
ti.grouped(y)
对张量 y 中的所有元素进行了一个打包操作
@ti.kernel
def array_op(x: ti.template(), y: ti.template()):
for I in ti.grouped(x):
# I is a vector of size x.dim() and data type i32
y[I] = I[0] + I[1]
如果 x 是二维张量,代码为:
@ti.kernel
def array_op(x: ti.template(), y: ti.template()):
for i, j in x:
y[i, j] = i + j
14.2 template 的反射
可在编译器获得 Tensor 大小的反射:
- 编译时可以获得 template 的大小
import taichi as ti
tensor = ti.var(ti.f32, shape=(4, 8, 16, 32, 64))
@ti.kernel
def print_tensor_size(x: ti.template()):
print(x.dim())
for i in ti.static(range(x.dim())):
print(x.shape()[i])
print_tensor_size(tensor)
14.3 编译期的分支语句
使用 ti.static(value)
让分支语句在编译期执行:
enable_projection = True
@ti.kernel
def static():
if ti.static(enable_projection): # No runtime overhead
x[0] = 1
14.4 编译期的循环展开
使用 ti.static(range(...))
进行循环展开:
import taichi as ti
ti.init()
x = ti.Vector(3, dt=ti.i32, shape=16)
@ti.kernel
def fill():
for i in x:
for j in ti.static(range(3)):
x[i][j] = j
print(x[i])
fill()
- 在向量/矩阵元素上循环,Taichi 的矩阵的索引必须是编译期常数,太极张量的下标可以是运行时变量:
eg:如果x是三维向量的一维张量,则作为x[tensor_index][matrix_index]
访问。第一个索引可以是变量,而第二个索引必须是常量 - 因此在循环张量时需要对其进行 unroll 操作
15. 变量别名
@ti.kernel
def my_kernel():
for i, j in tensor_a:
tensor_b[i, j] = some_function(tensor_a[i, j])
使用 ti.static()
对变量或者函数进行别名操作
@ti.kernel
def my_kernel():
a, b, fun = ti.static(tensor_a , tensor_b , some_function)
for i,j in a:
b[i,j] = fun(a[i,j])
16. 可微编程
正向编程求 f ( x ) f(\bold{x}) f(x) ,可微编程求 ∂ f ( x ) ∂ x \cfrac{\partial f(\bold{x})}{\partial \bold{x}} ∂x∂f(x)
- Taichi 支持反向模式自动微分(AutoDiff)
- 可以计算出反向传播 w.r.t.标量(损失)函数 f ( x ) f(\bold{x}) f(x) 的导数。
计算梯度的两种方式:
- 用
ti.Tape(loss)
,用于正向和梯度评估 - 显式使用 gradient kernels 进行梯度计算,运用上更多控制。
16.1 基于梯度的优化
m i n x L ( x ) = 1 2 ∑ i = 0 n − 1 ( x i − y i ) 2 min_{\bold{x}} \space \space L(\bold{x}) = \cfrac{1}{2} \sum_{i = 0}^{n-1}(\bold{x}_i - \bold{y}_i)^2 minx L(x)=21i=0∑n−1(xi−yi)2
16.1.1 指定参与梯度运算的变量
- 声明张量时分配梯度:
x = ti.var(dt=ti.f32, shape=n, needs_grad=True)
- 只有
needs_grad=True
才能求 L 关于 x 的导数 - 自变量和因变量都需要激活
16.1.2 定义损失函数
- 定义损失函数
kernel(s)
:
@ti.kernel
def reduce():
for i in range(n):
L[None] += 0.5 * (x[i] - y[i])**2
16.1.3 调用梯度值
- 定义梯度下降(在该处调用
x.grad[i]
):
@ti.kernel
def gradient_descent():
for i in x:
x[i] -= x.grad[i] * 0.1
- 也可用于时间积分(将
x.grad[i]
作为差分的数值):
@ti.kernel
def substep():
# Collide
......
# 显式时间
v[p] = (v[p] + ((-x.grad[p] / node_mass) + ti.Vector([0, -10, 0])) * dt) * math.exp(dt * -6)
x[p] += dt * v[p]
Note :
- 务必记得声明为
ti.kernel
16.1.4 使用微分器
- 对损失函数求导:
ti.Tape(loss=L): reduce()
# 使用100次梯度下降迭代进行优化
for k in range(100):
with ti.Tape(loss=L):
reduce()
print('Loss =', L[None])
gradient_descent()
16.1.5 代码示例
参考示例 :
ti example autodiff_minimization
ti example autodiff_regression
NOTE :
可以使用 rich
包:
pip install rich
这时可以使用语法高亮
ti example -p autodiff_minimization
可以将结果保存在当前目录下;
ti example -s autodiff_minimization
autodiff_minimization.py :
import taichi as ti
import random
n = 8
x = ti.field(dtype=ti.f32, shape=n, needs_grad=True)
y = ti.field(dtype=ti.f32, shape=n)
L = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
@ti.kernel
def reduce():
for i in range(n):
L[None] += 0.5 * (x[i] - y[i])**2
# Initialize vectors
for i in range(n):
x[i] = random.random()
y[i] = random.random()
@ti.kernel
def gradient_descent():
for i in x:
x[i] -= x.grad[i] * 0.1
# 使用100次梯度下降迭代进行优化
for k in range(100):
with ti.Tape(loss=L):
reduce()
print('Loss =', L[None])
gradient_descent()
for i in range(n):
# 现在你应该大约已经有 x[i] == y[i]
print(x[i], y[i])
损失函数:
L
=
1
2
(
x
i
−
y
i
)
2
L = \cfrac{1}{2} (x_i - y_i)^2
L=21(xi−yi)2
步长:0.1
16.2 自动微分的物理应用
16.2.1 应用1:势能梯度产生的力
势能对位移进行求导可以获得力:
f i = − ∂ U ( x ) ∂ x i \bold{f}_i = -\cfrac{\partial U(x)}{\partial \bold{x}_i} fi=−∂xi∂U(x)
手动导出梯度非常困难,使用 AutoDiff
:
-
分配一个零维张量来存储势能:
potential = ti.var(ti.f32, shape=())
-
定义从
x[i]
计算势能的正向核函数。 -
在
ti.Tape(loss=potential)
,调用正向核函数。 -
每个粒子上的力是
-x.grad[i]
。
16.2.2 应用2:对整个物理过程进行求导
Demo :DiffTaichi ( x t + 1 , v t + 1 , ⋯ ) = F ( x t , v t , ⋯ ) (\bold{x}_{t+1}, \bold{v}_{t+1}, \cdots) = \bold{F} (\bold{x}_{t}, \bold{v}_{t}, \cdots) (xt+1,vt+1,⋯)=F(xt,vt,⋯)
for k in range(100):
with ti.Tape(loss=loss):
for i in range(steps - 1):
simulate(i)
进行100次迭代,
在进行端到端的微分时,始终保持时间步长的整个计算存活周期。
- 即不只分配:
ti.Vector(3, dt=ti.f32, shape=(num_particles))
存储最新粒子; - 为整个模拟过程分配:
ti.Vector(3, dt=ti.f32, shape=(num_timesteps, num_particles))
。不要覆盖!
(使用检查点(本课程后面部分)来减少内存消耗。)
迭代过程中 x[i] -= x.grad[i] * alpha
,其中 alpha
是步长,如果使用最速下降或者共轭梯度等方法,步长则需要在迭代过程中进行计算。
17. 可视化
17.1 2D
gui = ti.GUI("Taichi MLS-MPM-128", res=512, background_color=0x112F41)
gui.circle/gui.circles(x.to_numpy(), radius=1.5, color=colors.to_numpy())
gui.line/triangle/set_image/show/...
17.2 3D
使用 ti.PLYWriter
导出三维粒子和网格
ti example export_ply/export_mesh
使用 Houdini
/ Blender
查看 3D 结果。