怎么用kernel.grad()计算N维输出对N维输入的grad

综合讨论区
#1

根据taichi文档里的示例改了一下，比如loss是2维的，x是16维的，在只计算一次func.grad的情况下可以得到一个2*16的Jacobian matrix吗？（因为用了grad_for,无法使用ti.ad.FwdMode）谢谢大家！

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=(2), needs_grad=True)

@ti.kernel
def func():
    for i in x:
       loss[0] += x[i] ** 2
       loss[1] += 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[0] = 1
loss.grad[1] = 1

func()
func.grad()

print(x.grad.shape)

# is it possible to get a 2 * 16 jacobian if func.grad() is only executed once?
#2

reverse mode autodiff一次只能算Jacobian matrix的一行，所以loss是2维的话，需要调用两次func.grad()，分别设置loss.grad[0] = 1, loss.grad[1] = 0和loss.grad[0] = 0, loss.grad[1] = 1，同时设置为1得到的是jacobian matrix两行的和

1 个赞
#3

谢谢解答！ 还有是打错了嘛，我测试了下是同时设置为1得到的是两行的和。
看了下这个视频，感觉一次只能计算一列或者一行是因为只存储了Tangents和Adjoints的值而没有记录产生这些值的function。如果记录了这些function，应该就能一次计算出整个Jacobian了。但可能考虑到内存和编译时间的原因就没存储这些function了？

#4

再补充一个踩到的小坑：如果用循环计算Jacobian matrix 一定记得要ti.ad.clear_all_gradients()

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=(2), needs_grad=True)

@ti.kernel
def func():
    for i in x:
       loss[0] += x[i] ** 2
       loss[1] += x[i]

for i in range(N):
    x[i] = i

for _ in range(10):
    loss.grad[0] = 1
    loss.grad[1] = 1
    func()
    func.grad()
    print(x.grad)
    #ti.ad.clear_all_gradients()
    # if not, grad will be accumulated
#5

是的是的，抱歉这里是我打错了，同时设置成1得到的是和

1 个赞
#6

是的，这里因为只开了jacobian matrix一行的内存，所以要清零一下gradient，不知道有没有看过这个例子taichi/jacobian.py at master · taichi-dev/taichi · GitHub

1 个赞
#7

明白了，谢谢解答！