Source code for pina.model.layers.residual

import torch
import torch.nn as nn
from ...utils import check_consistency


[docs] class ResidualBlock(nn.Module): """Residual block base class. Implementation of a residual block. .. seealso:: **Original reference**: He, Kaiming, et al. *Deep residual learning for image recognition.* Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.. DOI: `<https://arxiv.org/pdf/1512.03385.pdf>`_. """ def __init__( self, input_dim, output_dim, hidden_dim, spectral_norm=False, activation=torch.nn.ReLU(), ): """ Initializes the ResidualBlock module. :param int input_dim: Dimension of the input to pass to the feedforward linear layer. :param int output_dim: Dimension of the output from the residual layer. :param int hidden_dim: Hidden dimension for mapping the input (first block). :param bool spectral_norm: Apply spectral normalization to feedforward layers, defaults to False. :param torch.nn.Module activation: Cctivation function after first block. """ super().__init__() # check consistency check_consistency(spectral_norm, bool) check_consistency(input_dim, int) check_consistency(output_dim, int) check_consistency(hidden_dim, int) check_consistency(activation, torch.nn.Module) # assign variables self._spectral_norm = spectral_norm self._input_dim = input_dim self._output_dim = output_dim self._hidden_dim = hidden_dim self._activation = activation # create layers self._l1 = self._spect_norm(nn.Linear(input_dim, hidden_dim)) self._l2 = self._spect_norm(nn.Linear(hidden_dim, output_dim)) self._l3 = self._spect_norm(nn.Linear(input_dim, output_dim))
[docs] def forward(self, x): """Forward pass for residual block layer. :param torch.Tensor x: Input tensor for the residual layer. :return: Output tensor for the residual layer. :rtype: torch.Tensor """ y = self._activation(self._l1(x)) y = self._l2(y) x = self._l3(x) return y + x
def _spect_norm(self, x): """Perform spectral norm on the layers. :param x: A torch.nn.Module Linear layer :type x: torch.nn.Module :return: The spectral norm of the layer :rtype: torch.nn.Module """ return nn.utils.spectral_norm(x) if self._spectral_norm else x
import torch import torch.nn as nn
[docs] class EnhancedLinear(torch.nn.Module): """ A wrapper class for enhancing a linear layer with activation and/or dropout. :param layer: The linear layer to be enhanced. :type layer: torch.nn.Module :param activation: The activation function to be applied after the linear layer. :type activation: torch.nn.Module :param dropout: The dropout probability to be applied after the activation (if provided). :type dropout: float :Example: >>> linear_layer = torch.nn.Linear(10, 20) >>> activation = torch.nn.ReLU() >>> dropout_prob = 0.5 >>> enhanced_linear = EnhancedLinear(linear_layer, activation, dropout_prob) """ def __init__(self, layer, activation=None, dropout=None): """ Initializes the EnhancedLinear module. :param layer: The linear layer to be enhanced. :type layer: torch.nn.Module :param activation: The activation function to be applied after the linear layer. :type activation: torch.nn.Module :param dropout: The dropout probability to be applied after the activation (if provided). :type dropout: float """ super().__init__() # check consistency check_consistency(layer, nn.Module) if activation is not None: check_consistency(activation, nn.Module) if dropout is not None: check_consistency(dropout, float) # assign forward if (dropout is None) and (activation is None): self._model = torch.nn.Sequential(layer) elif (dropout is None) and (activation is not None): self._model = torch.nn.Sequential(layer, activation) elif (dropout is not None) and (activation is None): self._model = torch.nn.Sequential(layer, self._drop(dropout)) elif (dropout is not None) and (activation is not None): self._model = torch.nn.Sequential( layer, activation, self._drop(dropout) )
[docs] def forward(self, x): """ Forward pass through the enhanced linear module. :param x: Input tensor. :type x: torch.Tensor :return: Output tensor after passing through the enhanced linear module. :rtype: torch.Tensor """ return self._model(x)
def _drop(self, p): """ Applies dropout with probability p. :param p: Dropout probability. :type p: float :return: Dropout layer with the specified probability. :rtype: torch.nn.Dropout """ return torch.nn.Dropout(p)