Prepare for upcoming keras initializer change.
@@ -22,7 +22,9 @@ from acme.tf.networks import base
import sonnet as snt
import tensorflow as tf
-uniform_initializer = tf.initializers.VarianceScaling(
+
+def _uniform_initializer():
+ return tf.initializers.VarianceScaling(
distribution='uniform', mode='fan_out', scale=0.333)
@@ -40,10 +42,9 @@ class LayerNormMLP(snt.Module):
first layer and non-linearities (elu) on all but the last remaining layers.
"""
- def __init__(
- self,
+ def __init__(self,
layer_sizes: Sequence[int],
- w_init: Optional[snt.initializers.Initializer] = uniform_initializer,
+ w_init: Optional[snt.initializers.Initializer] = None,
activation: Callable[[tf.Tensor], tf.Tensor] = tf.nn.elu,
activate_final: bool = False):
"""Construct the MLP.
@@ -59,13 +60,13 @@ class LayerNormMLP(snt.Module):
super().__init__(name='feedforward_mlp_torso')
self._network = snt.Sequential([
- snt.Linear(layer_sizes[0], w_init=w_init),
+ snt.Linear(layer_sizes[0], w_init=w_init or _uniform_initializer()),
snt.LayerNorm(
axis=slice(1, None), create_scale=True, create_offset=True),
tf.nn.tanh,
snt.nets.MLP(
layer_sizes[1:],
- w_init=w_init,
+ w_init=w_init or _uniform_initializer(),
activation=activation,
activate_final=activate_final),
])
@@ -123,11 +124,12 @@ class LayerNormAndResidualMLP(snt.Module):
super().__init__(name='LayerNormAndResidualMLP')
# Create initial MLP layer.
- layers = [snt.nets.MLP([hidden_size], w_init=uniform_initializer)]
+ layers = [snt.nets.MLP([hidden_size], w_init=_uniform_initializer())]
# Follow it up with num_blocks MLPs with layernorm and residual connections.
for _ in range(num_blocks):
- mlp = snt.nets.MLP([hidden_size, hidden_size], w_init=uniform_initializer)
+ mlp = snt.nets.MLP([hidden_size, hidden_size],
+ w_init=_uniform_initializer())
layers.append(ResidualLayernormWrapper(mlp))
self._module = snt.Sequential(layers)
