可微分编程在科学计算与人工智能领域都有广泛的应用，自动微分（Automatic Differentiation）是一种对指定程序自动求导的技术，是实现可微分编程的重要工具。目前，Taichi Lang 支持反向模式自动微分，允许用户在 Taichi kernel 中编写可微分代码。

Taichi Lang 通过 Python scope 内的轻量级 tape 和 kernel 内部的源码转换 (Source Code Transformation) 来实现自动微分。tape 记录下启动的 Taichi kernel，源码转换在编译期间生成 gradient kernel，最后 tape 在反向传播期间以逆序重放 gradient kernel。
可微分物理仿真器是机器学习系统的常用组件，其提供的梯度信息可以使机器学习的收敛速度比无梯度算法（如 model-free reinforcement learning）快整整一个数量级，在机器人技能学习、形态优化、物理仿真学习等领域有重要应用。Taichi Lang 利用其原生基础架构生成高性能的梯度 kernel，这使得借此实现的自动微分功能十分适合进行物理仿真。
接下来，让我们一起动手实践：用 Taichi Lang 及其自动微分系统实现一个可微分物理仿真器。👇
目录
Contents
01 可微分光滑粒子流体力学仿真器
02 仿真器

03 神经网络

04 控制器

05 开始训练！
01 可微分光滑粒子流体力学仿真器 
Differentiable Smooth Particle Hydrodynamics Simulator
——
方法论
Methodology
光滑粒子流体动力学方法（SPH）是一种用于模拟连续介质力学的计算方法，多用于固体力学和流体力学的仿真。SPH 是一种基于粒子的仿真方法——即拉格朗日法，适合模拟喷泉等自由表面的流动。
控制与决策是人工智能领域的一类经典问题，我们可以尝试用可微分物理仿真来解决一个简单的控制问题。本篇文章中，作者将用 Taichi Lang 建一座“魔法”喷泉，喷泉应以最小的努力尽快击中指定目标。创建这座喷泉需要两个组件：一个是流体仿真器，一个是控制器。仿真器用来模拟水流的运动，控制器用来实现喷泉的“魔法”——驱动水流，而神经网络则用来拟合决定控制策略的函数。
02 仿真器
Simulator
——
流体的运动常常用纳维斯托克斯(N-S)方程来描述。流体仿真需要求解N-S方程，求解过程涉及压力与速度的解算。在该部分，我们将介绍一个基于弱可压光滑粒子流体动力学（WCSPH）的可微流体仿真器是如何逐步实现的。
可微分物理仿真器
Differentiable physical simulator
在不需要梯度信息的情况下，一个物理仿真器只缓存一个旧状态和一个新状态，这就是双缓存策略：物理仿真器每进行一次时间步的更新，两个缓存就会互换一次状态。基于链式法则计算导数需要物理量的所有历史状态，为了保留这些状态，仿真过程会被展开，即物理量每一仿真步的值都会被记录下来。
弱可压缩光滑粒子流体动力学
Weakly Compressible Smooth Particle Hydrodynamics
WCSPH 是 SPH 的变体。与严格保证液体不可压缩的常用投影方法不同，WCSPH 通过自定义状态方程（例如 Tait 方程）允许弱可压，避免了求解泊松方程耗时的问题。我们不会在这里深究 SPH 的原理。有兴趣的读者可以阅读 Eurographics Tutorial 2019 了解更多。
仿真展开
Unrolled simulation
由于整个仿真的历史都需要被存储下来，相关物理量分配内存时就需要增加一个维度 steps。 
pos = ti.Vector.field(3,
float
)

vel = ti.Vector.field(3,
float
)

acc = ti.Vector.field(3,
float
)

den = ti.field(
float
)

pre = ti.field(
float
)

ti.root.dense(ti.ijk, (batch_size, steps, particle_num)).place(pos, vel, acc, den, pre)
密度更新
Density update
粒子的密度由其相邻粒子的密度之和决定。由于本示例的粒子数量有限（小于 10k），为了简单起见，我们不对邻域搜寻过程应用任何加速算法（例如网格法、紧凑哈希），我们选择 poly6 作为核函数。
压力更新
Pressure update
基于 Tait 方程更新压力。为了避免出现颗粒密度不够的情况，低于静水密度的密度值会被统一拉齐到设定的下限。
压力和粘性力
Pressure and viscosity force
压力和粘性力的形式参见这篇论文👉🏻 Weakly compressible SPH for free surface flows ¹
时间积分
Time integration
时间积分采用半隐式欧拉方法。被水柱冲击的目标立方体的粒子速度始终设置为零。
仿真域的边界处理
Boundary handling 
为简单起见，边界的处理采用了基于碰撞的方法。流体粒子在与仿真域边界碰撞时会被推回仿真域，由系数 damping 控制碰撞过程中产生的速度损失。
03 神经网络
Neural network
——
得益于 Taichi Lang 的自动求导功能，利用其搭建的神经网络可以无缝嵌入同样基于 Taichi Lang 的物理仿真器。在动手实现控制器之前，我们先用 Taichi Lang 来搭建神经网络训练的流水线，这包括两个核心组件：神经网络和优化器。接下来，作者将向我们展示如何实现基于线性层（全连接层）的神经网络和随机梯度下降 (SGD) 优化器。
线性层
Linear layer
线性层由输入层、隐藏层和输出层组成。构建神经网络前，我们需要确定输入的维度。在该示例中，我们需要设置四个维度：第一个维度 n_models，用来同时验证多个训练好的模型，以提高训练后的验证效率；第二个维度 batching 可以提高网络训练效率；第三个维度 n_steps 用来记录每一步仿真相关的神经元的值；第四个维度 n_input代表神经网络输入特征。最终，我们的线性层的输入维度是 (n_models, batch_size, n_steps, n_input)。
示例中使用的神经网络可训练参数包含权重和偏差两类，权重有n_input * n_hidden 个参数，偏差有 n_hidden个参数。
weight、bias、hidden和 output的内存分配排布如下图所示：
*严格来说，线性层中不应该包含非线性激活函数，但是为了简化问题，我们的实现中包含了非线性激活函数。

随机梯度下降 
Stochastic Gradient Descent
随机梯度下降是当前流行的机器学习优化器。这里我们使用 Taichi Lang 可以很快实现一个 SGD。如下所示：
@ti.data_oriented

classSGD:
def__init__(self, params, lr)
:

self
.params = params

self
.lr = lr


defstep(self)
:

for
 w
inself
.
params:
self
._step(w)


    @ti.kernel

def_step(self, w: ti.template()
):

for
 I
in
 ti.grouped(w):

            w[I] -= min(max(w.grad[I], -
20.0
),
20.0
) *
self
.lr


defzero_grad(self)
:

for
 w
inself
.
params:
            w.grad.fill(
0
.
0
)
*请注意，如果仿真展开步数过长会造成梯度爆炸，我们应用了梯度裁剪来缓解这一问题。

04 控制器
Controller
——
在设置好仿真器和神经网络之后，我们需要搭建一个控制器，来决定如何驱动水流击中目标。如下图所示，在展开的每一步仿真中，我们都在 SPH 求解器前增加了一个控制器。控制器的输出作为输入的一部分送入求解器。
控制器由两个全连接层（fc1，fc2）组成。在全连接层 fc2 后增加了一个激活函数 tanh。控制器的输入是目标立方体的中心位置，控制器的输出是当前仿真步中施加到作用区域内驱动流体的力。作用区域被定义为一个立方块，位于仿真域的底部中心。驱动力会施加在作用区域内的每一个流体粒子上。为了加快训练过程的收敛，我们可以在仿真的前半部分添加一个小的向上的扰动力来引导梯度下降的方向。
05 开始训练！
Train it
——
数据集
Dataset
为了让控制器可以应对目标在各种不同位置的情况，我们生成了一个包含 80 个不同样本的训练数据集，每一个样本都是一个目标在三维空间内的位置坐标。我们将基于该数据集进行神经网络的训练。
网络和优化器的初始化
Network and Optimizer Initialization
接下来，我们用具有两个线性层的神经网络来构建控制器，并用来自线性层的可训练参数初始化基于 SGD 的优化器。
...

loss = ti.field(float, shape=(), needs_grad=
True
)

input_states = ti.field(float, shape=(model_num, steps, batch_size, n_input), needs_grad=
True
)


# Construct the fully-connected layers
fc1 = Linear(n_models=model_num,

             batch_size=batch_size,

             n_steps=steps, n_input=n_input, n_hidden=n_hidden,

             n_output=n_output, needs_grad=
True
, activation=
False
)

fc2 = Linear(n_models=model_num,

             batch_size=batch_size,

             n_steps=steps, n_input=n_output, n_hidden=n_hidden,

             n_output=n_output_act, needs_grad=
True
, activation=
True
)

fc1.weights_init()

fc2.weights_init()


# Feed trainable parameters to the optimizer
NNs = [fc1, fc2]

parameters = []

for
 layer
in
 NNs:

    parameters.extend(layer.parameters())

optimizer = SGD(params=parameters, lr=learning_rate)

...
损失函数
Loss function
该示例有两个目标，与之对应的损失函数包含损失项和正则化项。首先，我们的主要目标是喷泉击中指定物体，即流体粒子和物体的距离尽可能小。因此，损失项定义为所有仿真步骤中所有流体粒子与目标中心之间的最小距离。 
其次，我们需要让喷泉“以最小的努力”击中目标，即流体粒子在击中目标后不应偏离目标中心太远。因此，正则化项定义为所有仿真步骤中流体粒子与目标中心之间的最大距离。准备好所有组件后，我们就可以开始训练模型了。
“炼”起来！
Train
整个仿真包含 128 步，即喷泉需要在 128 步内击中目标。优化器会在一次完整的仿真后更新一遍神经网络参数。with ti.Tape(loss=loss)编译with scope 内的所有 Taichi kernel，生成相应的 gradient kernel 用于计算 loss 对可训练参数的导数。该模型可以在 10 次左右的优化迭代（每次的优化迭代包含 10 次仿真）中完成训练。
for
 opt_iter
in
 range(opt_iters):

    ...

for
 current_data_offset
in
 range(
0
, training_sample_num, batch_size):

        ...

with
 ti.Tape(loss=loss):

for
 i
in
 range(
1
, steps):

                initialize_density(i -
1
)

                update_density(i -
1
)

                update_pressure(i -
1
)

# Apply the NN based controller
                fc1.forward(i -
1
, input_states)

                fc2.forward(i -
1
, fc1.output)

                controller_output(i -
1
)

                apply_force(i -
1
)

                update_force(i -
1
)

                advance(i)

                boundary_handle(i)

                compute_dist(i)

            compute_loss(steps -
1
)

        optimizer.step()

        ...
#快来试试吧！
Play with the demo
——
Taichi 提供了带有训练模型的示例。如果你也想演示一下最终的效果，请务必安装最新版 Taichi：
pip install --upgrade taichi
运行 demo：
ti example diff_sph
本文介绍了基于 Taichi 自动微分实现的一个小例子，经过训练的喷泉可以用最小的努力击中目标。感兴趣的朋友可以阅读原文前往代码仓库体验 demo，也欢迎大家用 Taichi 写出可微分物理仿真或机器学习相关的应用并分享你的作品！
References
1. M. Becker and M. Teschner (2007). "Weakly compressible SPH for free surface flows". In: Proceedings of the 2007 ACM SIGGRAPH/Eurographics symposium on Computer animation. Eurographics Association, pp. 209–217
继续阅读
阅读原文