Source code for texar.torch.core.layers

# Copyright 2019 The Texar Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Various neural network layers
"""

import copy
import functools
from typing import Any, Callable, Dict, List, Optional, Type, Union

import torch
from torch import nn

from texar.torch.core import cell_wrappers as wrappers
from texar.torch.core.regularizers import L1L2, Regularizer
from texar.torch.hyperparams import HParams
from texar.torch.utils import utils
from texar.torch.utils.dtypes import is_str

__all__ = [
    'default_rnn_cell_hparams',
    'get_rnn_cell',
    'identity',
    'default_regularizer_hparams',
    'get_initializer',
    'get_regularizer',
    'get_activation_fn',
    'get_layer',
    'MaxReducePool1d',
    'AvgReducePool1d',
    'get_pooling_layer_hparams',
    'MergeLayer',
    'Flatten',
    'Identity',
]


[docs]def default_rnn_cell_hparams(): r"""Returns a `dict` of RNN cell hyperparameters and their default values. .. code-block:: python { "type": "LSTMCell", "input_size": 256, "kwargs": { "hidden_size": 256 }, "num_layers": 1, "dropout": { "input_keep_prob": 1.0, "output_keep_prob": 1.0, "state_keep_prob": 1.0, "variational_recurrent": False, }, "residual": False, "highway": False, } Here: `"type"`: str or cell class or cell instance The RNN cell type. This can be - The string name or full module path of a cell class. If class name is provided, the class must be in module :mod:`torch.nn.modules.rnn`, :mod:`texar.torch.core.cell_wrappers`, or :mod:`texar.torch.custom`. - A cell class. - An instance of a cell class. This is not valid if `"num_layers"` > 1. For example .. code-block:: python "type": "LSTMCell" # class name "type": "torch.nn.GRUCell" # module path "type": "my_module.MyCell" # module path "type": torch.nn.GRUCell # class "type": LSTMCell(hidden_size=100) # cell instance "type": MyCell(...) # cell instance `"kwargs"`: dict Keyword arguments for the constructor of the cell class. A cell is created by :python:`cell_class(**kwargs)`, where `cell_class` is specified in "type" above. Ignored if "type" is a cell instance. .. note:: It is unnecessary to specify `"input_size"` within `"kwargs"`. This value will be automatically filled based on layer index. .. note:: Although PyTorch uses `"hidden_size"` to denote the hidden layer size, we follow TensorFlow conventions and use `"num_units"`. `"num_layers"`: int Number of cell layers. Each layer is a cell created as above, with the same hyperparameters specified in `"kwargs"`. `"dropout"`: dict Dropout applied to the cell in **each** layer. See :class:`~texar.torch.core.cell_wrappers.DropoutWrapper` for details of the hyperparameters. If all `"\*_keep_prob"` = 1, no dropout is applied. Specifically, if `"variational_recurrent"` = `True`, the same dropout mask is applied across all time steps per batch. `"residual"`: bool If `True`, apply residual connection on the inputs and outputs of cell in **each** layer except the first layer. Ignored if `"num_layers"` = 1. `"highway"`: bool If True, apply highway connection on the inputs and outputs of cell in each layer except the first layer. Ignored if `"num_layers"` = 1. """ return { 'type': 'LSTMCell', 'kwargs': { 'num_units': 256, }, 'num_layers': 1, 'dropout': { 'input_keep_prob': 1.0, 'output_keep_prob': 1.0, 'state_keep_prob': 1.0, 'variational_recurrent': False, }, 'residual': False, 'highway': False, '@no_typecheck': ['type'] }
[docs]def default_regularizer_hparams(): r"""Returns the hyperparameters and their default values of a variable regularizer: .. code-block:: python { "type": "L1L2", "kwargs": { "l1": 0., "l2": 0. } } The default value corresponds to :class:`~texar.torch.core.regularizers.L1L2` and, with ``(l1=0, l2=0)``, disables regularization. """ return { "type": "L1L2", "kwargs": { "l1": 0., "l2": 0. } }
[docs]def get_rnn_cell(input_size, hparams=None): r"""Creates an RNN cell. See :func:`~texar.torch.core.default_rnn_cell_hparams` for all hyperparameters and default values. Args: input_size (int): Size of the input to the cell in the first layer. hparams (dict or HParams, optional): Cell hyperparameters. Missing hyperparameters are set to default values. Returns: A cell instance. Raises: ValueError: If ``hparams["num_layers"]``>1 and ``hparams["type"]`` is a class instance. """ if hparams is None or isinstance(hparams, dict): hparams = HParams(hparams, default_rnn_cell_hparams()) d_hp = hparams['dropout'] variational_recurrent = d_hp['variational_recurrent'] input_keep_prob = d_hp['input_keep_prob'] output_keep_prob = d_hp['output_keep_prob'] state_keep_prob = d_hp['state_keep_prob'] cells = [] num_layers = hparams['num_layers'] cell_kwargs = hparams['kwargs'].todict() # rename 'num_units' to 'hidden_size' following PyTorch conventions cell_kwargs['hidden_size'] = cell_kwargs['num_units'] del cell_kwargs['num_units'] for layer_i in range(num_layers): # Create the basic cell cell_type = hparams["type"] if layer_i == 0: cell_kwargs['input_size'] = input_size else: cell_kwargs['input_size'] = cell_kwargs['hidden_size'] if not isinstance(cell_type, str) and not isinstance(cell_type, type): if num_layers > 1: raise ValueError( "If 'num_layers'>1, then 'type' must be a cell class or " "its name/module path, rather than a cell instance.") cell_modules = ['texar.torch.core.cell_wrappers', # prefer our wrappers 'torch.nn.modules.rnn', 'texar.torch.custom'] cell = utils.check_or_get_instance(cell_type, cell_kwargs, cell_modules) if isinstance(cell, nn.RNNCellBase): cell = wrappers.wrap_builtin_cell(cell) # Optionally add dropout if (input_keep_prob < 1.0 or output_keep_prob < 1.0 or state_keep_prob < 1.0): # TODO: Would this result in non-final layer outputs being # dropped twice? cell = wrappers.DropoutWrapper( cell=cell, input_keep_prob=input_keep_prob, output_keep_prob=output_keep_prob, state_keep_prob=state_keep_prob, variational_recurrent=variational_recurrent) # Optionally add residual and highway connections if layer_i > 0: if hparams['residual']: cell = wrappers.ResidualWrapper(cell) if hparams['highway']: cell = wrappers.HighwayWrapper(cell) cells.append(cell) if hparams['num_layers'] > 1: cell = wrappers.MultiRNNCell(cells) else: cell = cells[0] return cell
def identity(inputs: torch.Tensor): r"""Returns a tensor with the same content as the input tensor. Arguments: inputs: The input tensor. Returns: A tensor of the same shape, type, and content. """ return inputs
[docs]def get_regularizer(hparams=None): r"""Returns a variable regularizer instance. See :func:`~texar.torch.core.default_regularizer_hparams` for all hyperparameters and default values. The "type" field can be a subclass of :class:`~texar.torch.core.regularizers.Regularizer`, its string name or module path, or a class instance. Args: hparams (dict or HParams, optional): Hyperparameters. Missing hyperparameters are set to default values. Returns: A :class:`~texar.torch.core.regularizers.Regularizer` instance. `None` if :attr:`hparams` is `None` or taking the default hyperparameter value. Raises: ValueError: The resulting regularizer is not an instance of :class:`~texar.torch.core.regularizers.Regularizer`. """ if hparams is None: return None if isinstance(hparams, dict): hparams = HParams(hparams, default_regularizer_hparams()) rgl = utils.check_or_get_instance( hparams.type, hparams.kwargs.todict(), ["texar.torch.core.regularizers", "texar.torch.custom"]) if not isinstance(rgl, Regularizer): raise ValueError("The regularizer must be an instance of " "texar.torch.core.regularizers.Regularizer.") if isinstance(rgl, L1L2) and rgl.l1 == 0. and rgl.l2 == 0.: return None return rgl
[docs]def get_initializer(hparams=None) \ -> Optional[Callable[[torch.Tensor], torch.Tensor]]: r"""Returns an initializer instance. Args: hparams (dict or HParams, optional): Hyperparameters with the structure .. code-block:: python { "type": "initializer_class_or_function", "kwargs": { # ... } } The `"type"` field can be a function name or module path. If name is provided, it be must be from one the following modules: :torch_docs:`torch.nn.init <nn.html#torch-nn-init>` and :mod:`texar.torch.custom`. Besides, the `"type"` field can also be an initialization function called with :python:`initialization_fn(**kwargs)`. In this case `"type"` can be the function, or its name or module path. If no keyword argument is required, `"kwargs"` can be omitted. Returns: An initializer instance. `None` if :attr:`hparams` is `None`. """ if hparams is None: return None kwargs = hparams.get('kwargs', {}) if isinstance(kwargs, HParams): kwargs = kwargs.todict() modules = ['torch.nn.init', 'torch', 'texar.torch.custom'] initializer_fn = utils.get_function(hparams['type'], modules) initializer = functools.partial(initializer_fn, **kwargs) return initializer
[docs]def get_activation_fn(fn_name: Optional[Union[str, Callable[[torch.Tensor], torch.Tensor]]] = None, kwargs: Union[HParams, Dict, None] = None) \ -> Optional[Callable[[torch.Tensor], torch.Tensor]]: r"""Returns an activation function `fn` with the signature `output = fn(input)`. If the function specified by :attr:`fn_name` has more than one arguments without default values, then all these arguments except the input feature argument must be specified in :attr:`kwargs`. Arguments with default values can also be specified in :attr:`kwargs` to take values other than the defaults. In this case a partial function is returned with the above signature. Args: fn_name (str or callable): An activation function, or its name or module path. The function can be: - Built-in function defined in :torch_docs:`torch.nn.functional<nn.html#torch-nn-functional>` - User-defined activation functions in module :mod:`texar.torch.custom`. - External activation functions. Must provide the full module path, e.g., ``"my_module.my_activation_fn"``. kwargs (optional): A `dict` or instance of :class:`~texar.torch.HParams` containing the keyword arguments of the activation function. Returns: An activation function. `None` if :attr:`fn_name` is `None`. """ if fn_name is None: return None fn_modules = ['torch', 'torch.nn.functional', 'texar.torch.custom', 'texar.torch.core.layers'] activation_fn_ = utils.get_function(fn_name, fn_modules) activation_fn = activation_fn_ # Make a partial function if necessary if kwargs is not None: if isinstance(kwargs, HParams): kwargs = kwargs.todict() def _partial_fn(features): return activation_fn_(features, **kwargs) activation_fn = _partial_fn return activation_fn
[docs]def get_layer(hparams: Union[HParams, Dict[str, Any]]) -> nn.Module: r"""Makes a layer instance. The layer must be an instance of :torch_nn:`Module`. Args: hparams (dict or HParams): Hyperparameters of the layer, with structure: .. code-block:: python { "type": "LayerClass", "kwargs": { # Keyword arguments of the layer class # ... } } Here: `"type"`: str or layer class or layer instance The layer type. This can be - The string name or full module path of a layer class. If the class name is provided, the class must be in module :torch_nn:`Module`, :mod:`texar.torch.core`, or :mod:`texar.torch.custom`. - A layer class. - An instance of a layer class. For example .. code-block:: python "type": "Conv1D" # class name "type": "texar.torch.core.MaxReducePooling1D" # module path "type": "my_module.MyLayer" # module path "type": torch.nn.Module.Linear # class "type": Conv1D(filters=10, kernel_size=2) # cell instance "type": MyLayer(...) # cell instance `"kwargs"`: dict A dictionary of keyword arguments for constructor of the layer class. Ignored if :attr:`"type"` is a layer instance. - Arguments named "activation" can be a callable, or a `str` of the name or module path to the activation function. - Arguments named "\*_regularizer" and "\*_initializer" can be a class instance, or a `dict` of hyperparameters of respective regularizers and initializers. See - Arguments named "\*_constraint" can be a callable, or a `str` of the name or full path to the constraint function. Returns: A layer instance. If ``hparams["type"]`` is a layer instance, returns it directly. Raises: ValueError: If :attr:`hparams` is `None`. ValueError: If the resulting layer is not an instance of :torch_nn:`Module`. """ if hparams is None: raise ValueError("`hparams` must not be `None`.") layer_type = hparams["type"] if not is_str(layer_type) and not isinstance(layer_type, type): layer = layer_type else: layer_modules = ["torch.nn", "texar.torch.core", "texar.torch.custom"] layer_class: Type[nn.Module] = utils.check_or_get_class( layer_type, layer_modules) if isinstance(hparams, dict): if (layer_class.__name__ == "Linear" and "in_features" not in hparams["kwargs"]): raise ValueError("\"in_features\" should be specified for " "\"torch.nn.{}\"".format(layer_class.__name__)) elif (layer_class.__name__ in ["Conv1d", "Conv2d", "Conv3d"] and "in_channels" not in hparams["kwargs"]): raise ValueError("\"in_channels\" should be specified for " "\"torch.nn.{}\"".format(layer_class.__name__)) default_kwargs: Dict[str, Any] = {} default_hparams = {"type": layer_type, "kwargs": default_kwargs} hparams = HParams(hparams, default_hparams) # this case needs to be handled separately because nn.Sequential # does not accept kwargs if layer_type == "Sequential": names: List[str] = [] layer = nn.Sequential() sub_hparams = hparams.kwargs.layers for hparam in sub_hparams: sub_layer = get_layer(hparam) name = utils.uniquify_str(sub_layer.__class__.__name__, names) names.append(name) layer.add_module(name=name, module=sub_layer) else: layer = utils.get_instance(layer_type, hparams.kwargs.todict(), layer_modules) if not isinstance(layer, nn.Module): raise ValueError("layer must be an instance of `torch.nn.Module`.") return layer
[docs]class MaxReducePool1d(nn.Module): r"""A subclass of :torch_nn:`Module`. Max Pool layer for 1D inputs. The same as :torch_nn:`MaxPool1d` except that the pooling dimension is entirely reduced (i.e., `pool_size=input_length`). """
[docs] def forward(self, # type: ignore input: torch.Tensor) -> torch.Tensor: output, _ = torch.max(input, dim=2) return output
[docs]class AvgReducePool1d(nn.Module): r"""A subclass of :torch_nn:`Module`. Avg Pool layer for 1D inputs. The same as :torch_nn:`AvgPool1d` except that the pooling dimension is entirely reduced (i.e., `pool_size=input_length`). """
[docs] def forward(self, # type: ignore input: torch.Tensor) -> torch.Tensor: return torch.mean(input, dim=2)
_POOLING_TO_REDUCE = { "MaxPool1d": "MaxReducePool1d", "AvgPool1d": "AvgReducePool1d", torch.nn.MaxPool1d: MaxReducePool1d, torch.nn.AvgPool1d: AvgReducePool1d }
[docs]def get_pooling_layer_hparams(hparams: Union[HParams, Dict[str, Any]]) \ -> Dict[str, Any]: r"""Creates pooling layer hyperparameters `dict` for :func:`get_layer`. If the :attr:`hparams` sets `'pool_size'` to `None`, the layer will be changed to the respective reduce-pooling layer. For example, :torch_docs:`torch.conv.MaxPool1d <nn.html#torch.nn.Conv1d>` is replaced with :class:`~texar.torch.core.MaxReducePool1d`. """ if isinstance(hparams, HParams): hparams = hparams.todict() new_hparams = copy.copy(hparams) kwargs = new_hparams.get('kwargs', None) if kwargs and kwargs.get('kernel_size', None) is None: pool_type = hparams['type'] new_hparams['type'] = _POOLING_TO_REDUCE.get(pool_type, pool_type) kwargs.pop('kernel_size', None) kwargs.pop('stride', None) kwargs.pop('padding', None) return new_hparams
[docs]class MergeLayer(nn.Module): r"""A subclass of :torch_nn:`Module`. A layer that consists of multiple layers in parallel. Input is fed to each of the parallel layers, and the outputs are merged with a specified mode. Args: layers (list, optional): A list of :torch_docs:`torch.nn.Module <nn.html#module>` instances, or a list of hyperparameter dictionaries each of which specifies `"type"` and `"kwargs"` of each layer (see the `hparams` argument of :func:`get_layer`). If `None`, this layer degenerates to a merging operator that merges inputs directly. mode (str): Mode of the merge op. This can be: - :attr:`'concat'`: Concatenates layer outputs along one dim. Tensors must have the same shape except for the dimension specified in `dim`, which can have different sizes. - :attr:`'elemwise_sum'`: Outputs element-wise sum. - :attr:`'elemwise_mul'`: Outputs element-wise product. - :attr:`'sum'`: Computes the sum of layer outputs along the dimension given by `dim`. For example, given `dim=1`, two tensors of shape `[a, b]` and `[a, c]` respectively will result in a merged tensor of shape `[a]`. - :attr:`'mean'`: Computes the mean of layer outputs along the dimension given in `dim`. - :attr:`'prod'`: Computes the product of layer outputs along the dimension given in `dim`. - :attr:`'max'`: Computes the maximum of layer outputs along the dimension given in `dim`. - :attr:`'min'`: Computes the minimum of layer outputs along the dimension given in `dim`. - :attr:`'and'`: Computes the `logical and` of layer outputs along the dimension given in `dim`. - :attr:`'or'`: Computes the `logical or` of layer outputs along the dimension given in `dim`. - :attr:`'logsumexp'`: Computes log(sum(exp(elements across the dimension of layer outputs))) dim (int): The dim to use in merging. Ignored in modes :attr:`'elemwise_sum'` and :attr:`'elemwise_mul'`. """ _functions: Dict[str, Callable[[torch.Tensor, int], torch.Tensor]] = { "sum": torch.sum, "mean": torch.mean, "prod": torch.prod, "max": lambda tensors, dim: torch.max(tensors, dim)[0], "min": lambda tensors, dim: torch.min(tensors, dim)[0], "and": torch.all, "or": torch.any, "logsumexp": torch.logsumexp } def __init__(self, layers: Optional[List[nn.Module]] = None, mode: str = 'concat', dim: Optional[int] = None): super().__init__() self._mode = mode self._dim = dim self._layers: Optional[nn.ModuleList] = None if layers is not None: if len(layers) == 0: raise ValueError( "'layers' must be either None or a non-empty list.") self._layers = nn.ModuleList() for layer in layers: if isinstance(layer, nn.Module): self._layers.append(layer) else: self._layers.append(get_layer(hparams=layer))
[docs] def forward(self, input: torch.Tensor) -> torch.Tensor: # type: ignore r"""Feed input to every containing layer and merge the outputs. Args: input: The input tensor. Returns: The merged tensor. """ layer_outputs: List[torch.Tensor] if self._layers is None: layer_outputs = input if not isinstance(layer_outputs, (list, tuple)): layer_outputs = [layer_outputs] else: layer_outputs = [] for layer in self._layers: layer_output = layer(input) layer_outputs.append(layer_output) # the merge dimension cannot be determined until we get the output from # individual layers. # In case of reduce pooling operations, feature dim is removed and # channel dim is merged. # In non-reduce pooling operations, feature dim is merged. dim = self._dim if self._dim is not None else -1 if self._mode == 'concat': outputs = torch.cat(tensors=layer_outputs, dim=dim) elif self._mode == 'elemwise_sum': outputs = layer_outputs[0] for i in range(1, len(layer_outputs)): outputs = torch.add(outputs, layer_outputs[i]) elif self._mode == 'elemwise_mul': outputs = layer_outputs[0] for i in range(1, len(layer_outputs)): outputs = torch.mul(outputs, layer_outputs[i]) elif self._mode in self._functions: _concat = torch.cat(tensors=layer_outputs, dim=dim) outputs = self._functions[self._mode](_concat, dim) else: raise ValueError("Unknown merge mode: '%s'" % self._mode) return outputs
@property def layers(self) -> Optional[nn.ModuleList]: r"""The list of parallel layers. """ return self._layers
[docs]class Flatten(nn.Module): r"""Flatten layer to flatten a tensor after convolution.""" def forward(self, # type: ignore input: torch.Tensor) -> torch.Tensor: return input.view(input.size()[0], -1)
[docs]class Identity(nn.Module): r"""Identity activation layer.""" def forward(self, # type: ignore input: torch.Tensor) -> torch.Tensor: return input