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