我复制并简化了这段极其简单的代码,只是为了开始使用神经切线和 jax
from jax import random
import neural_tangents as nt
init_fn, apply_fn, kernel_fn = nt.stax.serial(
nt.stax.Dense(512), nt.stax.Relu(),
nt.stax.Dense(1)
)
key1, _ = random.split(random.PRNGKey(1))
x = random.normal(key1, (1, 1)) # training features
y = random.uniform(key1, shape=(1, 1)) # training targets
predict_fn = nt.predict.gradient_descent_mse_ensemble(kernel_fn, x, y)
predictions = predict_fn(x, get='ntk')
print(y)
print("********")
print(predictions)
据我了解,这训练了一个“无限宽”的神经网络,以将一个训练示例实数拟合到所需的目标(也是单个实数)。
所以y
和predictions
应该是一样的。我正在训练一个例子,我有一个无限强大的模型,在我的脑海中这两个应该是相同的。他们不是。它打印:
[[0.12462091]]
********
[[[[-0.45066586]]]]
更重要的是,如果我将训练示例的数量更改为 3,那么现在两个打印的尺寸不匹配!我希望找到两个向量,每个向量包含三个数字。我得到的是:
[[0.6283829 ]
[0.53268707]
[0.89792836]]
********
[[[[0.0453045 ]
[0.01134446]
[0.17562252]]]
[[[0.01215818]
[0.00304447]
[0.0471311 ]]]
[[[0.1367114 ]
[0.03423319]
[0.5299607 ]]]]
很明显,我有一个致命的误解。文档没有帮助我。任何人都可以阐明这个问题吗?