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 长度相同的输出序列。