Differentiable Programming¶
约 209 个字 37 行代码 预计阅读时间 1 分钟
Taichi 通过ti.ad.Tape()
或更灵活的kernel.grad()
语法来实现对梯度的计算
使用 ti.ad.Tape() ¶
Example
在上面的代码示例中,有以下几个点需要关注
- 对于要计算梯度的变量,在声明的时候需要设置一个参数
needs_grad=True
,才能通过with ti.ad.Tape()
语句自动微分
Tip
每个变量在声明的时候都要写这个参数会很麻烦,可以通过 Taichi 的 API ti.root.lazy_grad()
来实现
ti.ad.Tape()
中的参数必须是 0 维的 field
使用 kernel.grad() ¶
和使用ti.ad.Tape()
不同,在使用kernel.grad()
之前需要把输出变量的梯度值手动设置为 1
Example
import taichi as ti
ti.init()
N = 16
x = ti.field(dtype=ti.f32, shape=N, needs_grad=True)
loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
loss2 = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
@ti.kernel
def func():
for i in x:
loss[None] += x[i] ** 2
loss2[None] += x[i]
for i in range(N):
x[i] = i
# Set the `grad` of the output variables to `1` before calling `func.grad()`.
loss.grad[None] = 1
loss2.grad[None] = 1
func()
func.grad()
for i in range(N):
assert x.grad[i] == i * 2 + 1
在使用反向的kernel.grad()
之前应该先运行前向的kernel()
,避免全局的 Taichi field 被修改了但是没有做前向运算,导致梯度错误