1

我复制并简化了这段极其简单的代码,只是为了开始使用神经切线和 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)

据我了解,这训练了一个“无限宽”的神经网络,以将一个训练示例实数拟合到所需的目标(也是单个实数)。

所以ypredictions应该是一样的。我正在训练一个例子,我有一个无限强大的模型,在我的脑海中这两个应该是相同的。他们不是。它打印:

[[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 ]]]]

很明显,我有一个致命的误解。文档没有帮助我。任何人都可以阐明这个问题吗?

4

0 回答 0