如何解决了解梯度函数中的 JAX argnums 参数
我正在尝试了解 JAX 梯度函数中 argnums
的行为。
假设我有以下功能:
def make_mse(x,t):
def mse(w,b):
return np.sum(jnp.power(x.dot(w) + b - t,2))/2
return mse
我采用以下方式进行渐变:
w_gradient,b_gradient = grad(make_mse(train_data,y),(0,1))(w,b)
argnums= (0,1)
在这种情况下,但它是什么意思?关于哪些变量计算梯度?如果我改用 argnums=0
会有什么不同?
另外,我可以使用相同的函数来获得 Hessian 矩阵吗?
我查看了有关它的 JAX help 部分,但无法弄清楚
解决方法
当您将多个 argnum 传递给 grad 时,结果是一个返回梯度元组的函数,相当于您分别计算了每个:
def f(x,y):
return x ** 2 + x * y + y ** 2
df_dxy = grad(f,argnums=(0,1))
df_dx = grad(f,argnums=0)
df_dy = grad(f,argnums=1)
x = 3.0
y = 4.25
assert df_dxy(x,y) == (df_dx(x,y),df_dy(x,y))
如果要计算混合二阶导数,可以通过重复应用梯度来实现:
d2f_dxdy = grad(grad(f,argnums=0),argnums=1)
assert d2f_dxdy(x,y) == 1
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。