感謝您閱讀我的問題！
我剛剛學習了 Jax 中的自定義 grad 函式，我發現 JAX 定義自定義函式所采用的方法非常優雅。
不過有一件事困擾著我。
我創建了一個包裝器，使松散卷積看起來像 PyTorch conv2d。
from jax import numpy as jnp
from jax.random import PRNGKey, normal
from jax import lax
from torch.nn.modules.utils import _ntuple
import jax
from jax.nn.initializers import normal
from jax import grad
torch_dims = {0: ('NC', 'OI', 'NC'), 1: ('NCH', 'OIH', 'NCH'), 2: ('NCHW', 'OIHW', 'NCHW'), 3: ('NCHWD', 'OIHWD', 'NCHWD')}
def conv(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
n = len(input.shape) - 2
if type(stride) == int:
stride = _ntuple(n)(stride)
if type(padding) == int:
padding = [(i, i) for i in _ntuple(n)(padding)]
if type(dilation) == int:
dilation = _ntuple(n)(dilation)
return lax.conv_general_dilated(lhs=input, rhs=weight, window_strides=stride, padding=padding, lhs_dilation=dilation, rhs_dilation=None, dimension_numbers=torch_dims[n], feature_group_count=1, batch_group_count=1, precision=None, preferred_element_type=None)
問題是我找不到使用其 grad 函式的方法：
init = normal()
rng = PRNGKey(42)
x = init(rng, [128, 3, 224, 224])
k = init(rng, [64, 3, 3, 3])
y = conv(x, k)
grad(conv)(y, k)
這就是我得到的。
ValueError: conv_general_dilated lhs feature dimension size divided by feature_group_count must equal the rhs input feature dimension size, but 64 // 1 != 3.
請幫忙！
uj5u.com熱心網友回復：
當我使用最新版本的 jax 和 jaxlib ( jax==0.2.22; jaxlib==0.1.72)運行您的代碼時，我看到以下錯誤：
TypeError: Gradient only defined for scalar-output functions. Output had shape: (128, 64, 222, 222).
如果我創建一個使用 的標量輸出函式，conv梯度似乎有效：
result = grad(lambda x, k: conv(x, k).sum())(x, k)
print(result.shape)
# (128, 3, 224, 224)
如果您使用的是舊版本的 JAX，您可能會嘗試更新到更新的版本 - 您看到的錯誤可能是由于已修復的錯誤所致。
轉載請註明出處，本文鏈接：https://www.uj5u.com/net/317857.html