jax.example_libraries.stax 模块#
Stax 是一个从头开始构建的，小巧但灵活的神经网络规范库。
你可能并不想导入这个模块！Stax 仅用作示例库。JAX 有许多其他功能更完善的神经网络库，包括 Google 的 Flax 和 DeepMind 的 Haiku。
- jax.example_libraries.stax.AvgPool(window_shape, strides=None, padding='VALID', spec=None)[源代码]#
用于池化层的层构建函数。
- jax.example_libraries.stax.BatchNorm(axis=(0, 1, 2), epsilon=1e-05, center=True, scale=True, beta_init=<function zeros>, gamma_init=<function ones>)[源代码]#
用于批归一化层的层构建函数。
- jax.example_libraries.stax.Conv(out_chan, filter_shape, strides=None, padding='VALID', W_init=None, b_init=<function normal.<locals>.init>)#
用于通用卷积层的层构建函数。
- jax.example_libraries.stax.Conv1DTranspose(out_chan, filter_shape, strides=None, padding='VALID', W_init=None, b_init=<function normal.<locals>.init>)#
用于通用转置卷积层的层构建函数。
- jax.example_libraries.stax.ConvTranspose(out_chan, filter_shape, strides=None, padding='VALID', W_init=None, b_init=<function normal.<locals>.init>)#
用于通用转置卷积层的层构建函数。
- jax.example_libraries.stax.Dense(out_dim, W_init=<function variance_scaling.<locals>.init>, b_init=<function normal.<locals>.init>)[源代码]#
用于密集（全连接）层的层构造函数。
- jax.example_libraries.stax.GeneralConv(dimension_numbers, out_chan, filter_shape, strides=None, padding='VALID', W_init=None, b_init=<function normal.<locals>.init>)[源代码]#
用于通用卷积层的层构建函数。
- jax.example_libraries.stax.GeneralConvTranspose(dimension_numbers, out_chan, filter_shape, strides=None, padding='VALID', W_init=None, b_init=<function normal.<locals>.init>)[源代码]#
用于通用转置卷积层的层构建函数。
- jax.example_libraries.stax.MaxPool(window_shape, strides=None, padding='VALID', spec=None)[源代码]#
用于池化层的层构建函数。
- jax.example_libraries.stax.SumPool(window_shape, strides=None, padding='VALID', spec=None)[源代码]#
用于池化层的层构建函数。
- jax.example_libraries.stax.parallel(*layers)[源代码]#
用于并行组合层的组合器。
此组合器生成的层通常与 FanOut 和 FanInSum 层一起使用。
- 参数：
*layers – 一系列层，每个层都是一个 (init_fun, apply_fun) 对。
- 返回值：
一个新的层，表示一个 (init_fun, apply_fun) 对，表示给定层序列的并行组合。 特别是，返回的层接收一系列输入并返回与参数 layers 长度相同的一系列输出。