def _build_dense_output_layer(cell_output_size: int,
hparams: HParams) -> Optional[nn.Sequential]:
r"""Build the output layers.
Args:
cell_output_size: The output size of the rnn cell.
hparams (dict or HParams): Hyperparameters. Missing hyperparameters
will be set to default values. See
:meth:`default_hparams` for the hyperparameter structure and
default values.
Returns:
A :torch_nn:`Sequential` module containing the output layers.
"""
nlayers = hparams.num_layers
if nlayers <= 0:
return None
layer_size = _to_list(hparams.layer_size, 'output_layer.layer_size',
nlayers)
dropout_layer_ids = _to_list(hparams.dropout_layer_ids)
other_kwargs = hparams.other_dense_kwargs or {}
if isinstance(other_kwargs, HParams):
other_kwargs = other_kwargs.todict()
if not isinstance(other_kwargs, dict):
raise ValueError(
"hparams 'output_layer.other_dense_kwargs' must be a dict.")
output_layers: List[nn.Module] = []
for i in range(nlayers):
if i in dropout_layer_ids:
# TODO: Variational dropout is not implemented.
output_layers.append(nn.Dropout(p=hparams.dropout_rate))
dense_layer = nn.Linear(
in_features=(cell_output_size if i == 0 else layer_size[i - 1]),
out_features=layer_size[i],
**other_kwargs)
output_layers.append(dense_layer)
if i == nlayers - 1:
activation = hparams.final_layer_activation
else:
activation = hparams.activation
if activation is not None:
layer_hparams = {"type": activation, "kwargs": {}}
activation_layer = layers.get_layer(hparams=layer_hparams)
output_layers.append(activation_layer)
if nlayers in dropout_layer_ids:
output_layers.append(nn.Dropout(p=hparams.dropout_rate))
return nn.Sequential(*output_layers)
def _forward_output_layers(inputs,
input_size,
output_layer,
time_major,
hparams,
mode,
sequence_length=None):
"""Forwards inputs through the output layers.
Args:
inputs: A Tensor of shape `[batch_size, max_time] + input_size` if
:attr:`time_major=False`, or shape
`[max_time, batch_size] + input_size` if :attr:`time_major=True`.
Returns:
A pair :attr:`(outputs, outputs_size), where
- :attr:`outputs`: A Tensor of shape \
`[batch_size, max_time] + outputs_size`.
- :attr:`outputs_size`: An `int` or 1D `int` array representing the \
output size.
"""
if output_layer is None:
return inputs, input_size
if hparams is None:
# output_layer was passed in from the constructor
if isinstance(output_layer, (list, tuple)):
raise ValueError('output_layer must not be a list or tuple.')
output, output_size = _forward_single_output_layer(
inputs, input_size, output_layer)
else:
# output_layer was built based on hparams
output_layer = _to_list(output_layer)
dropout_layer_ids = _to_list(hparams.dropout_layer_ids)
if len(dropout_layer_ids) > 0:
training = is_train_mode(mode)
output = inputs
output_size = input_size
for i, layer in enumerate(output_layer):
if i in dropout_layer_ids:
output = _apply_dropout(output, time_major, hparams, training)
output, output_size = _forward_single_output_layer(
output, output_size, layer)
if len(output_layer) in dropout_layer_ids:
output = _apply_dropout(output, time_major, hparams, training)
if sequence_length is not None:
output = mask_sequences(output,
sequence_length,
time_major=time_major,
tensor_rank=3)
return output, output_size
def _build_dense_output_layer(hparams):
nlayers = hparams.num_layers
if nlayers <= 0:
return None
layer_size = _to_list(
hparams.layer_size, 'output_layer.layer_size', nlayers)
other_kwargs = hparams.other_dense_kwargs or {}
if isinstance(other_kwargs, HParams):
other_kwargs = other_kwargs.todict()
if not isinstance(other_kwargs, dict):
raise ValueError(
"hparams 'output_layer.other_dense_kwargs' must be a dict.")
dense_layers = []
for i in range(nlayers):
if i == nlayers - 1:
activation = hparams.final_layer_activation
else:
activation = hparams.activation
kwargs_i = {"units": layer_size[i],
"activation": activation,
"name": "dense_%d" % (i+1)}
kwargs_i.update(other_kwargs)
layer_hparams = {"type": "Dense", "kwargs": kwargs_i}
dense_layers.append(layers.get_layer(hparams=layer_hparams))
if len(dense_layers) == 1:
dense_layers = dense_layers[0]
return dense_layers