credit.models.graph =================== .. py:module:: credit.models.graph Attributes ---------- .. autoapisummary:: credit.models.graph.n_variables Classes ------- .. autoapisummary:: credit.models.graph.GraphResTransfGRU credit.models.graph.GraphResidualBlock credit.models.graph.TransformerConv credit.models.graph.LayerNorm credit.models.graph.GateCell Functions --------- .. autoapisummary:: credit.models.graph.apply_spectral_norm Module Contents --------------- .. py:function:: apply_spectral_norm(model) .. py:class:: GraphResTransfGRU(n_variables=4, n_surface_variables=7, n_static_variables=3, levels=15, hidden_size=128, dim_head=32, dropout=0, n_blocks=3, history_len=2, edge_path='/glade/derecho/scratch/dgagne/credit_scalers/grid_edge_pairs_125.nc', use_spectral_norm=True, use_edge_attr=True) Bases: :py:obj:`credit.models.base_model.BaseModel` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: n_variables :value: 4 .. py:attribute:: n_static_variables :value: 3 .. py:attribute:: n_surface_variables :value: 7 .. py:attribute:: histroy_len :value: 2 .. py:attribute:: n_levels :value: 15 .. py:attribute:: state_vars :value: 67 .. py:attribute:: total_n_vars :value: 140 .. py:attribute:: n_blocks :value: 3 .. py:attribute:: hidden_size :value: 128 .. py:attribute:: dim_head :value: 32 .. py:attribute:: heads :value: 4 .. py:attribute:: dropout :value: 0 .. py:attribute:: edge_path :value: '/glade/derecho/scratch/dgagne/credit_scalers/grid_edge_pairs_125.nc' .. py:attribute:: use_spectral_norm :value: True .. py:attribute:: use_edge_attr :value: True .. py:attribute:: encoder .. py:attribute:: decoder .. py:attribute:: graph_blocks .. py:attribute:: gated_unit .. py:attribute:: graph_norm_layers .. py:attribute:: enc_norm .. py:attribute:: dec_norm .. py:method:: forward(x) .. py:method:: load_graph() .. py:method:: to(*args, **kwargs) Move and/or cast the parameters and buffers. This can be called as .. function:: to(device=None, dtype=None, non_blocking=False) :noindex: .. function:: to(dtype, non_blocking=False) :noindex: .. function:: to(tensor, non_blocking=False) :noindex: .. function:: to(memory_format=torch.channels_last) :noindex: Its signature is similar to :meth:`torch.Tensor.to`, but only accepts floating point or complex :attr:`dtype`\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:`dtype` (if given). The integral parameters and buffers will be moved :attr:`device`, if that is given, but with dtypes unchanged. When :attr:`non_blocking` is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices. See below for examples. .. note:: This method modifies the module in-place. :param device: the desired device of the parameters and buffers in this module :type device: :class:`torch.device` :param dtype: the desired floating point or complex dtype of the parameters and buffers in this module :type dtype: :class:`torch.dtype` :param tensor: Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module :type tensor: torch.Tensor :param memory_format: the desired memory format for 4D parameters and buffers in this module (keyword only argument) :type memory_format: :class:`torch.memory_format` :returns: self :rtype: Module Examples:: >>> # xdoctest: +IGNORE_WANT("non-deterministic") >>> linear = nn.Linear(2, 2) >>> linear.weight Parameter containing: tensor([[ 0.1913, -0.3420], [-0.5113, -0.2325]]) >>> linear.to(torch.double) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1913, -0.3420], [-0.5113, -0.2325]], dtype=torch.float64) >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1) >>> gpu1 = torch.device("cuda:1") >>> linear.to(gpu1, dtype=torch.half, non_blocking=True) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1914, -0.3420], [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1') >>> cpu = torch.device("cpu") >>> linear.to(cpu) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1914, -0.3420], [-0.5112, -0.2324]], dtype=torch.float16) >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble) >>> linear.weight Parameter containing: tensor([[ 0.3741+0.j, 0.2382+0.j], [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128) >>> linear(torch.ones(3, 2, dtype=torch.cdouble)) tensor([[0.6122+0.j, 0.1150+0.j], [0.6122+0.j, 0.1150+0.j], [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128) .. py:class:: GraphResidualBlock(input_size, hidden_size, out_size, heads=1, dropout=0, use_spectral_norm=True, use_edge_attr=True) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: input_size .. py:attribute:: hidden_size .. py:attribute:: out_size .. py:attribute:: heads :value: 1 .. py:attribute:: dropout :value: 0 .. py:attribute:: use_spectral_norm :value: True .. py:attribute:: use_edge_attr :value: True .. py:attribute:: transformer .. py:attribute:: physics_linear .. py:attribute:: merge_linear .. py:attribute:: norm_layer .. py:method:: forward(x, edge_index, edge_attr) .. py:class:: TransformerConv(in_channels: Union[int, Tuple[int, int]], out_channels: int, heads: int = 1, concat: bool = True, beta: bool = False, dropout: float = 0.0, edge_dim: Optional[int] = None, bias: bool = True, root_weight: bool = True, **kwargs) Bases: :py:obj:`torch_geometric.nn.conv.MessagePassing` Adapted from https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/conv/transformer_conv.html#TransformerConv To added additional `batch` dimension to the input since the graph doesn't change instead of using PyG's graph_batch which will duplicate the same graph. .. py:attribute:: _alpha :type: torch_geometric.typing.OptTensor .. py:attribute:: in_channels .. py:attribute:: out_channels .. py:attribute:: heads :value: 1 .. py:attribute:: beta :value: False .. py:attribute:: root_weight :value: True .. py:attribute:: concat :value: True .. py:attribute:: dropout :value: 0.0 .. py:attribute:: edge_dim :value: None .. py:attribute:: lin_key .. py:attribute:: lin_query .. py:attribute:: lin_value .. py:method:: reset_parameters() Resets all learnable parameters of the module. .. py:method:: forward(x: Union[torch.Tensor, torch_geometric.typing.PairTensor], edge_index: torch_geometric.typing.Adj, edge_attr: torch_geometric.typing.OptTensor = None, return_attention_weights: Optional[bool] = None) -> Union[torch.Tensor, Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], Tuple[torch.Tensor, torch_geometric.typing.SparseTensor]] Runs the forward pass of the module. .. py:method:: message(query_i: torch.Tensor, key_j: torch.Tensor, value_j: torch.Tensor, edge_attr: torch_geometric.typing.OptTensor, index: torch.Tensor, ptr: torch_geometric.typing.OptTensor, size_i: Optional[int]) -> torch.Tensor Constructs messages from node :math:`j` to node :math:`i` in analogy to :math:`\phi_{\mathbf{\Theta}}` for each edge in :obj:`edge_index`. This function can take any argument as input which was initially passed to :meth:`propagate`. Furthermore, tensors passed to :meth:`propagate` can be mapped to the respective nodes :math:`i` and :math:`j` by appending :obj:`_i` or :obj:`_j` to the variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`. .. py:method:: __repr__() -> str .. py:class:: LayerNorm(in_channels, eps=1e-05) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: in_channels .. py:attribute:: eps :value: 1e-05 .. py:attribute:: weight .. py:attribute:: bias .. py:method:: reset_parameters() Resets all learnable parameters of the module. .. py:method:: forward(x) .. py:class:: GateCell(hidden_size) Bases: :py:obj:`torch.nn.Module` Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool .. py:attribute:: hidden_size .. py:attribute:: z_x_ln .. py:attribute:: z_h_ln .. py:attribute:: r_x_ln .. py:attribute:: r_h_ln .. py:attribute:: h_x_ln .. py:attribute:: h_h_ln .. py:method:: forward(x, h) .. py:data:: n_variables :value: 4