【GAMES201学习笔记】00 - Taichi 编程

参考文档

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.kernelti.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 程序的几个阶段

  1. 初始化: ti.init(...)
  2. Tensor 定义(到计算阶段就无法继续定义): ti.var, ti.Vector, ti.Matrix
  3. 计算阶段 (启动 kernels)
  4. 可选: 重新启动太极系统(清除内存,销毁所有变量和 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/r23

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}} xf(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=0n1(xiyi)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(xiyi)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=xiU(x)

手动导出梯度非常困难,使用 AutoDiff

  1. 分配一个零维张量来存储势能: potential = ti.var(ti.f32, shape=())

  2. 定义从 x[i] 计算势能的正向核函数。

  3. ti.Tape(loss=potential) ,调用正向核函数。

  4. 每个粒子上的力是 -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

《Taichi 中文文档》 GUI 系统

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

《Taichi 中文文档》 导出 PLY 文件

使用 ti.PLYWriter 导出三维粒子和网格

ti example export_ply/export_mesh

使用 Houdini / Blender 查看 3D 结果。

  • 17
    点赞
  • 35
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值