Neural Tangents实战：10个核心函数详解与代码示例

Neural Tangents实战：10个核心函数详解与代码示例

【免费下载链接】neural-tangentsFast and Easy Infinite Neural Networks in Python项目地址: https://gitcode.com/gh_mirrors/ne/neural-tangents

Neural Tangents是一个强大的Python库，专注于提供快速且简单的无限神经网络实现。本文将深入解析该库中10个核心函数，帮助新手和普通用户快速掌握其使用方法，轻松构建和分析无限宽度神经网络模型。

1. empirical_ntk_fn：计算神经网络的神经正切核

empirical_ntk_fn是Neural Tangents库中最核心的函数之一，用于计算神经网络的神经正切核（NTK）。神经正切核在研究无限宽度神经网络的行为中起着关键作用。

def empirical_ntk_fn( f: ApplyFn, params: PyTree, *, trace_axes: Axes = (), diagonal_axes: Axes = (), vmap_axes: VMapAxes = 0, implementation: str = 'auto', batch_size: Optional[int] = None, fwd: bool = True, reverse: bool = True, device: Optional[jax.Device] = None, precision: Optional[jax.lax.Precision] = None ) -> EmpiricalGetKernelFn:

该函数位于neural_tangents/_src/empirical.py，通过计算神经网络参数的雅可比矩阵来构建核函数，可用于分析网络的泛化能力和训练动态。

2. linearize：线性化神经网络

linearize函数用于在给定参数处线性化神经网络，这对于研究神经网络的局部行为和构建简化模型非常有用。

def linearize(f: ApplyFn, params: PyTree) -> ApplyFn: def f_lin(p, *args, **kwargs): return f(params, *args, **kwargs) + jax.jvp( lambda p_: f(p_, *args, **kwargs), (params,), (tree_sub(p, params),) )[1] return f_lin

这个函数在neural_tangents/_src/empirical.py中定义，返回一个新的函数，该函数在参数params附近线性近似原函数f。

3. taylor_expand：泰勒展开神经网络

taylor_expand函数提供了比简单线性化更高级的近似方法，通过泰勒级数展开来近似神经网络。

def taylor_expand(f: ApplyFn, params: PyTree, degree: int) -> ApplyFn: def taylorize_r(f, params, dparams, degree, current_degree): if current_degree == degree: def f_jvp(p): return f(p) return f_jvp else: pushfwd = jax.jvp(f, (params,), (dparams,))[1] return taylorize_r(pushfwd, params, dparams, degree, current_degree + 1) def f_tayl(p, *args, **kwargs): dparams = tree_sub(p, params) taylor = f(params, *args, **kwargs) for d in range(1, degree + 1): term = taylorize_r(lambda p_: f(p_, *args, **kwargs), params, dparams, d, 1)(params) taylor = tree_add(taylor, tree_div(term, factorial(d))) return taylor return f_tayl

位于neural_tangents/_src/empirical.py的此函数允许用户指定展开的阶数，提供更精确的非线性近似。

4. empirical_nngp_fn：计算神经网络高斯过程核

empirical_nngp_fn函数用于计算神经网络的神经预测高斯过程（NNGP）核，这是另一种描述无限宽度神经网络行为的重要工具。

def empirical_nngp_fn( f: ApplyFn, params: PyTree, *, trace_axes: Axes = (), diagonal_axes: Axes = (), vmap_axes: VMapAxes = 0, batch_size: Optional[int] = None, device: Optional[jax.Device] = None, precision: Optional[jax.lax.Precision] = None ) -> EmpiricalGetKernelFn:

这个函数在neural_tangents/_src/empirical.py中实现，与NTK不同，NNGP核描述了无限宽度神经网络输出之间的协方差结构。

5. Dense：构建全连接层

Dense函数是构建神经网络的基础模块，用于创建全连接层。

def Dense( out_dim: int, W_init: Initializer = ntk_initializers.glorot_normal(), b_init: Initializer = ntk_initializers.zeros, parameterization: str = 'ntk', use_bias: bool = True, precision: Optional[jax.lax.Precision] = None ) -> InternalLayer:

在neural_tangents/_src/stax/linear.py中定义的这个函数提供了灵活的全连接层实现，支持不同的参数化方式和初始化方法。

6. Conv：构建卷积层

Conv函数用于创建卷积层，这是处理图像数据的关键组件。

def Conv( features: int, kernel_size: Sequence[int], strides: Sequence[int] = (1, 1), padding: Union[str, Sequence[Union[str, int]]] = 'SAME', W_init: Initializer = ntk_initializers.glorot_normal(), b_init: Initializer = ntk_initializers.zeros, parameterization: str = 'ntk', use_bias: bool = True, precision: Optional[jax.lax.Precision] = None ) -> InternalLayer:

位于neural_tangents/_src/stax/linear.py的Conv函数支持多种卷积参数设置，包括 kernel 大小、步幅和填充方式。

7. monte_carlo_kernel_fn：蒙特卡洛核函数估计

monte_carlo_kernel_fn函数提供了一种通过蒙特卡洛采样来估计核函数的方法，特别适用于复杂模型。

def monte_carlo_kernel_fn( f: ApplyFn, kernel_fn: GetKernelFn[Kernel], params: PyTree, *, n_samples: Union[int, Sequence[int]] = 10, key: jnp.ndarray, batch_size: Optional[int] = None, vmap_axes: VMapAxes = 0, implementation: str = 'auto', trace_axes: Axes = (), diagonal_axes: Axes = (), device: Optional[jax.Device] = None, precision: Optional[jax.lax.Precision] = None ) -> EmpiricalGetKernelFn:

在neural_tangents/_src/monte_carlo.py中实现的这个函数通过多次采样来近似计算核函数，为分析大型复杂网络提供了实用工具。

8. empirical_ntk_vp_fn：神经正切核向量乘积

empirical_ntk_vp_fn函数计算神经正切核与向量的乘积，这对于高效计算某些优化和推理任务非常重要。

def empirical_ntk_vp_fn( f: ApplyFn, params: PyTree, *, vmap_axes: VMapAxes = 0, implementation: str = 'auto', batch_size: Optional[int] = None, device: Optional[jax.Device] = None, precision: Optional[jax.lax.Precision] = None ) -> NTKVPFn: def ntk_vp_fn(cotangents: PyTree) -> PyTree: # 实现细节... return ntk_vp_fn

这个函数在neural_tangents/_src/empirical.py中定义，提供了一种高效计算NTK向量乘积的方法，避免了显式构建大型核矩阵。

9. Identity：恒等层

Identity函数创建一个恒等映射层，在构建复杂网络结构时非常有用。

def Identity() -> InternalLayer: """Layer that returns its inputs unchanged.""" init_fn = lambda rng, input_shape: ((), input_shape) apply_fn = lambda params, inputs, **kwargs: inputs kernel_fn = lambda k, **kwargs: k mask_fn = lambda mask, input_shape: mask return init_fn, apply_fn, kernel_fn, mask_fn

位于neural_tangents/_src/stax/linear.py的这个简单但重要的函数在构建残差连接等网络结构时不可或缺。

10. AvgPool：平均池化层

AvgPool函数实现平均池化操作，用于降采样和特征提取。

def AvgPool( window_shape: Sequence[int], strides: Optional[Sequence[int]] = None, padding: Union[str, Sequence[Union[str, int]]] = 'SAME' ) -> InternalLayer: return _Pool( window_shape, strides, padding, reduce_fn=lax.reduce_mean, kernel_reduce_fn=lambda x, **kwargs: x.mean(**kwargs) )

在neural_tangents/_src/stax/linear.py中定义的这个函数提供了平均池化功能，是构建卷积神经网络的重要组件。

如何开始使用Neural Tangents

要开始使用Neural Tangents，首先需要克隆仓库：

git clone https://gitcode.com/gh_mirrors/ne/neural-tangents

然后可以参考examples/目录中的示例代码，如examples/empirical_ntk.py和examples/infinite_fcn.py，了解如何使用这些核心函数构建和分析无限宽度神经网络。

Neural Tangents库提供了丰富的工具来探索无限宽度神经网络的特性，通过掌握这些核心函数，您可以更深入地理解神经网络的行为，并构建更强大的机器学习模型。无论是研究还是应用开发，这些工具都能为您提供有价值的 insights 和高效的实现方式。

【免费下载链接】neural-tangentsFast and Easy Infinite Neural Networks in Python项目地址: https://gitcode.com/gh_mirrors/ne/neural-tangents