关于自动求导的中间变量

看到fem99.py这个代码里面的自动求导部分,有点困惑,其中的中间变量如D_i, phi_i等,它们不是taichi的域,也不会被追踪导数,那么链式法则在这里是如何生效的呢?taichi的编译器会自动替换掉中间变量嘛?

pos = ti.Vector.field(2, float, NV, needs_grad=True)
vel = ti.Vector.field(2, float, NV)
f2v = ti.Vector.field(3, int, NF)  # ids of three vertices of each face
B = ti.Matrix.field(2, 2, float, NF)
F = ti.Matrix.field(2, 2, float, NF, needs_grad=True)
V = ti.field(float, NF)
phi = ti.field(float, NF)  # potential energy of each face (Neo-Hookean)
U = ti.field(float, (), needs_grad=True)  # total potential energy
@ti.kernel
def update_U():
    for i in range(NF):
        ia, ib, ic = f2v[i]
        a, b, c = pos[ia], pos[ib], pos[ic]
        V[i] = abs((a - c).cross(b - c))
        D_i = ti.Matrix.cols([a - c, b - c])
        F[i] = D_i @ B[i]
    for i in range(NF):
        F_i = F[i]
        log_J_i = ti.log(F_i.determinant())
        phi_i = mu / 2 * ((F_i.transpose() @ F_i).trace() - 2)
        phi_i -= mu * log_J_i
        phi_i += lam / 2 * log_J_i**2
        phi[i] = phi_i
        U[None] += V[i] * phi_i

Hi @Sky, 非常欢迎来到太极论坛。

如果你想要了解DIFFTAICHI的 具体设计和实现,可以看herehere

1 个赞

谢谢~

我查看了一下提供的paper,看到了这一段

不过由于对编译器的术语不大熟悉,不是很确定我的理解是否准确

所以如果从用户的角度来说的话,定义一些非太极域的中间变量不会有影响,只需要链式法则上涉及的太极域变量设置为

needs_grad=True

就没问题了对吗?

嗯嗯,是的。

除了needs_grad=True之外,你还可以设置ti.root.lazy_grad()

可以看Github上关于autodiff的讨论:here