credit.models.graph

credit.models.graph#

Attributes#

n_variables

Classes#

`GraphResTransfGRU`	Base class for all neural network modules.
`GraphResidualBlock`	Base class for all neural network modules.
`TransformerConv`	Adapted from
`LayerNorm`	Base class for all neural network modules.
`GateCell`	Base class for all neural network modules.

Functions#

apply_spectral_norm(model)

Module Contents#

credit.models.graph.apply_spectral_norm(model)#

class credit.models.graph.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: 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 to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

n_variables = 4#

n_static_variables = 3#

n_surface_variables = 7#

histroy_len = 2#

n_levels = 15#

state_vars = 67#

total_n_vars = 140#

n_blocks = 3#

hidden_size = 128#

dim_head = 32#

heads = 4#

dropout = 0#

edge_path = '/glade/derecho/scratch/dgagne/credit_scalers/grid_edge_pairs_125.nc'#

use_spectral_norm = True#

use_edge_attr = True#

encoder#

decoder#

graph_blocks#

gated_unit#

graph_norm_layers#

enc_norm#

dec_norm#

forward(x)#

load_graph()#

to(*args, **kwargs)#

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When 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.

Parameters:

device (torch.device) – the desired device of the parameters and buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

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)

class credit.models.graph.GraphResidualBlock(input_size, hidden_size, out_size, heads=1, dropout=0, use_spectral_norm=True, use_edge_attr=True)#

Bases: 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 to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

input_size#

hidden_size#

out_size#

heads = 1#

dropout = 0#

use_spectral_norm = True#

use_edge_attr = True#

transformer#

physics_linear#

merge_linear#

norm_layer#

forward(x, edge_index, edge_attr)#

class credit.models.graph.TransformerConv(in_channels: int | Tuple[int, int], out_channels: int, heads: int = 1, concat: bool = True, beta: bool = False, dropout: float = 0.0, edge_dim: int | None = None, bias: bool = True, root_weight: bool = True, **kwargs)#

Bases: 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.

_alpha: torch_geometric.typing.OptTensor#

in_channels#

out_channels#

heads = 1#

beta = False#

root_weight = True#

concat = True#

dropout = 0.0#

edge_dim = None#

lin_key#

lin_query#

lin_value#

reset_parameters()#: Resets all learnable parameters of the module.

forward(x: torch.Tensor | torch_geometric.typing.PairTensor, edge_index: torch_geometric.typing.Adj, edge_attr: torch_geometric.typing.OptTensor = None, return_attention_weights: bool | None = None) → torch.Tensor | Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] | Tuple[torch.Tensor, torch_geometric.typing.SparseTensor]#: Runs the forward pass of the module.

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: int | None) → torch.Tensor#: Constructs messages from node \(j\) to node \(i\) in analogy to \(\phi_{\mathbf{\Theta}}\) for each edge in edge_index. This function can take any argument as input which was initially passed to propagate(). Furthermore, tensors passed to propagate() can be mapped to the respective nodes \(i\) and \(j\) by appending _i or _j to the variable name, .e.g. x_i and x_j.

__repr__() → str#

class credit.models.graph.LayerNorm(in_channels, eps=1e-05)#

Bases: 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 to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

in_channels#

eps = 1e-05#

weight#

bias#

reset_parameters()#: Resets all learnable parameters of the module.

forward(x)#

class credit.models.graph.GateCell(hidden_size)#

Bases: 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 to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

hidden_size#

z_x_ln#

z_h_ln#

r_x_ln#

r_h_ln#

h_x_ln#

h_h_ln#

forward(x, h)#

credit.models.graph.n_variables = 4#