在使用ti.root.lazy_grad()的时候，经常遇到的一个报错是:
RuntimeError: Gradients of loss are not allocated, please use ti.field(…, needs_grad=True) for all fields that are required by autodiff.
一开始我真是十分疑惑啊，因为使用的field都是通过ti.root.dense或者ti.root.pointer放进去的，好像没办法加上needs_grad = True. 经过一番探索以后，发现上面这个报错主要是下面三个原因导致的。
-
所有与loss相关的计算必须是一个ti.kernel或者ti.func。有的计算因为不需要并行，所以一开始没加ti.kernel修饰，然而这么做就求不到导了（这点Taichi文档里说明了的）。
-
对于field的赋值一定要在ti.root.lazy_grad()之后。比如下面的示例中如果对v[None]的赋值在lazy_grad之前也会报上面的错。
import taichi as ti
ti.init()
real = ti.f32
scalar = lambda: ti.field(dtype=real)
vec = lambda: ti.Vector.field(3, dtype=real)
target = ti.Vector([0.5, 0.5, 0.5])
n = 2 # particle number
x, v, x_avg = vec(), vec(), vec()
loss = scalar()
ti.root.place(v, loss, x_avg)
ti.root.dense(ti.i, n).place(x)
#v[None] = [2.0, 1.0, 3.0]
ti.root.lazy_grad()
v[None] = [2.0, 1.0, 3.0]
@ti.kernel
def compute_x_avg():
for p in x:
x_avg[None] += 1 / n * x[p]
@ti.kernel
def compute_loss():
dist = x_avg[None] - target
loss[None] = (dist[0] ** 2 + dist[1] ** 2 + dist[2] ** 2) ** 0.5
@ti.kernel
def substep():
for p in x:
x[p] += v[None]
with ti.ad.Tape(loss):
substep()
compute_x_avg()
compute_loss()
print('dloss/dv =', v.grad)
-
对于不参与loss计算但又必须在ti.root.lazy_root()前初始化的值（比如传给一个class用于loss计算的常数值），可以用ti.ndarray()储存，就不会影响lazy_root了。
-
有的时候发现grad的值都是0。一种可能的情况是在tape里没有用atomic add改变gloabal variable的值。比如下面的代码中，如果直接用=号给v赋值且ini_v是在tape中进行的，那么v.grad都会是0。
import taichi as ti
ti.init()
real = ti.f32
scalar = lambda: ti.field(dtype=real)
vec = lambda: ti.Vector.field(3, dtype=real)
target = ti.Vector([0.5, 0.5, 0.5])
n = 2 # particle number
x, v, x_avg = vec(), vec(), vec()
loss = scalar()
ti.root.place(loss, x_avg)
ti.root.dense(ti.i, n).place(x,v)
ti.root.lazy_grad()
@ti.kernel
def ini_v():
for p in v:
v[p] = [0.1, 0.2, 0.3]
#v[p] += [0.1, 0.2, 0.3] # if ini_v() is inside tape, use atomic add so the v.grad will be non-zero
@ti.kernel
def compute_x_avg():
for p in x:
x_avg[None] += 1 / n * x[p]
@ti.kernel
def compute_loss():
dist = x_avg[None] - target
loss[None] = (dist[0] ** 2 + dist[1] ** 2 + dist[2] ** 2) ** 0.5
@ti.kernel
def substep():
for p in x:
x[p] += v[p]
ini_v()
with ti.ad.Tape(loss):
#ini_v() # v.grad will become 0!
substep()
compute_x_avg()
compute_loss()
print('dloss/dv =', v.grad)
- 如果grad的值是nan，可能是因为反向传播时除了个0。（虽然前向传播的时候看起来一切正常）。在所有涉及除号的地方加一个很小的值，如下，世界就变得美好了起来。
g_v = grid_v_out[base + offset] / (grid_m[base + offset] + 1e-10)
请问大家第2点和第4点是因为什么而导致的呢？
私货时间：diffTaichi 里提到check pointing 有两种方式，一个是 D1. RECOMPUTATION WITHIN TIME STEPS，另一个是D2. SEGMENT-WISE RECOMPUTATION。examples里找到了D1.的实现，想问有实现D2的例子吗？（懒惰的我）