python - 了解其梯度函数中的 JAX argnums 参数
问题描述
我试图了解argnums
JAX 渐变函数中的行为。假设我有以下功能:
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 帮助部分,但无法弄清楚
解决方案
当您将多个 argnums 传递给 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
推荐阅读
- android - 当应用程序在 android webRTC/AppRTC 中处于后台时停止传输音频
- java - 我的位置正在更新,而我没有使用 altbeacon android studio 移动
- flutter - 如何在列表视图构建器中显示特定索引中的项目
- powershell - kubectl 运行中的格式错误的命令
- qt - KeyPress 未在 Android 平板电脑上触发
- javascript - 如何使用 map area="rect" 突出显示图像的可点击区域 - Ionic 4
- javascript - While 循环元素状态 cypress
- laravel - Laravel 从另一个来源获取值
- c# - 表中的 Jquery 自动完成
- c# - 如何自定义渲染器?