credit.models

Submodules

• credit.models.base_model
• credit.models.checkpoint
• credit.models.crossformer
• credit.models.crossformer_diffusion
• credit.models.crossformer_ensemble
• credit.models.debugger_model
• credit.models.fuxi
• credit.models.graph
• credit.models.reset
• credit.models.swin
• credit.models.unet
• credit.models.unet404
• credit.models.unet_attention_modules
• credit.models.unet_diffusion

Attributes

logger
model_types

Classes

CrossFormer	Base class for all neural network modules.
SegmentationModel	Base class for all neural network modules.
SegmentationModel404	Base class for all neural network modules.
Fuxi
SwinTransformerV2Cr	Swin Transformer V2
GraphResTransfGRU	Base class for all neural network modules.
DebuggerModel	Base class for all neural network modules.
CrossFormerWithNoise	CrossFormer variant with pixel-wise noise injection in both encoder and decoder stages.
CrossFormerDiffusion	Base class for all neural network modules.
UnetDiffusion	Base class for all neural network modules.
ModifiedGaussianDiffusion	Base class for all neural network modules.

Functions

load_fsdp_or_checkpoint_policy(conf)
load_model(conf[, load_weights, model_name])
load_model_name(conf, model_name[, load_weights])

Package Contents

class credit.models.CrossFormer(image_height: int = 640, patch_height: int = 1, image_width: int = 1280, patch_width: int = 1, frames: int = 2, channels: int = 4, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, levels: int = 15, dim: tuple = (64, 128, 256, 512), depth: tuple = (2, 2, 8, 2), dim_head: int = 32, global_window_size: tuple = (5, 5, 2, 1), local_window_size: int = 10, cross_embed_kernel_sizes: tuple = ((4, 8, 16, 32), (2, 4), (2, 4), (2, 4)), cross_embed_strides: tuple = (4, 2, 2, 2), attn_dropout: float = 0.0, ff_dropout: float = 0.0, use_spectral_norm: bool = True, attention_type: str = None, interp: bool = True, upsample_v_conv: bool = False, padding_conf: dict = None, post_conf: dict = None, **kwargs)

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.

image_height = 640

image_width = 1280

patch_height = 1

patch_width = 1

upsample_v_conv = False

frames = 2

channels = 4

surface_channels = 7

levels = 15

use_spectral_norm = True

use_interp = True

use_padding

use_post_block

input_only_channels = 3

input_channels = 70

output_channels = 67

layers

cube_embedding

up_block1

up_block2

up_block3

forward(x)

rk4(x)

class credit.models.SegmentationModel(image_height=640, image_width=1280, frames=2, channels=4, surface_channels=7, input_only_channels=3, output_only_channels=0, levels=16, rk4_integration=False, architecture=None, post_conf=None, **kwargs)

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.

image_height = 640

image_width = 1280

frames = 2

channels = 4

surface_channels = 7

levels = 16

rk4_integration = False

model

use_post_block

forward(x)

rk4(x)

class credit.models.SegmentationModel404(conf)

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.

variables

frames

static_variables

model

forward(x)

class credit.models.Fuxi(image_height=640, patch_height=16, image_width=1280, patch_width=16, levels=15, frames=2, frame_patch_size=2, dim=1536, num_groups=32, channels=4, surface_channels=7, input_only_channels=0, output_only_channels=0, num_heads=8, depth=48, window_size=7, use_spectral_norm=True, interp=True, proj_drop=0, attn_drop=0, drop_path=0, padding_conf=None, post_conf=None, **kwargs)

Bases: credit.models.base_model.BaseModel

Parameters:

• img_size (Sequence[int], optional) – T, Lat, Lon.

• patch_size (Sequence[int], optional) – T, Lat, Lon.

• in_chans (int, optional) – number of input channels.

• out_chans (int, optional) – number of output channels.

• dim (int, optional) – number of embed channels.

• num_groups (Sequence[int] | int, optional) – number of groups to separate the channels into.

• num_heads (int, optional) – Number of attention heads.

• window_size (int | tuple[int], optional) – Local window size.

use_interp = True

use_spectral_norm = True

use_padding

use_post_block

cube_embedding

u_transformer

fc

patch_size

input_resolution

out_chans = 67

img_size

channels = 4

surface_channels = 7

levels = 15

forward(x: torch.Tensor)

class credit.models.SwinTransformerV2Cr(img_size: Tuple[int, int] = (224, 224), patch_size: int = 4, window_size: int | None = None, img_window_ratio: int = 32, channels: int = 4, levels: int = 15, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, frames: int = 1, embed_dim: int = 96, depths: Tuple[int, Ellipsis] = (2, 2, 6, 2), num_heads: Tuple[int, Ellipsis] = (3, 6, 12, 24), mlp_ratio: float = 4.0, init_values: float | None = 0.0, drop_rate: float = 0.0, proj_drop_rate: float = 0.0, attn_drop_rate: float = 0.0, drop_path_rate: float = 0.0, norm_layer: Type[torch.nn.Module] = nn.LayerNorm, extra_norm_period: int = 0, extra_norm_stage: bool = False, sequential_attn: bool = False, global_pool: str = 'avg', weight_init='skip', full_pos_embed: bool = False, rel_pos: bool = True, checkpoint_stages: bool = False, residual: bool = False, use_spectral_norm: bool = False, padding_conf: dict = None, post_conf: dict = None, **kwargs: Any)

Bases: credit.models.base_model.BaseModel

Swin Transformer V2

A PyTorch impl ofSwin Transformer V2: Scaling Up Capacity and Resolution -

https://arxiv.org/pdf/2111.09883

Parameters:

• img_size – Input resolution.

• window_size – Window size. If None, img_size // window_div

• img_window_ratio – Window size to image size ratio.

• patch_size – Patch size.

• in_chans – Number of input channels.

• depths – Depth of the stage (number of layers).

• num_heads – Number of attention heads to be utilized.

• embed_dim – Patch embedding dimension.

• num_classes – Number of output classes.

• mlp_ratio – Ratio of the hidden dimension in the FFN to the input channels.

• drop_rate – Dropout rate.

• proj_drop_rate – Projection dropout rate.

• attn_drop_rate – Dropout rate of attention map.

• drop_path_rate – Stochastic depth rate.

• norm_layer – Type of normalization layer to be utilized.

• extra_norm_period – Insert extra norm layer on main branch every N (period) blocks in stage

• extra_norm_stage – End each stage with an extra norm layer in main branch

• sequential_attn – If true sequential self-attention is performed.

• padding_conf (dict) – padding configuration

• post_conf (dict) – configuration for postblock processing

use_padding

patch_size: int = 4

img_size: Tuple[int, int] = (224, 224)

window_size: int

num_features: int = 96

frames = 1

in_chans = 70

out_chans = 67

feature_info = []

full_pos_embed = False

checkpoint_stages = False

residual = False

depth

use_post_block

patch_embed

stages

head

use_spectral_norm = False

forward_features(x: torch.Tensor) → torch.Tensor

forward_head(x: torch.Tensor) → torch.Tensor

forward(x: torch.Tensor) → torch.Tensor

update_input_size(new_img_size: Tuple[int, int] | None = None, new_window_size: int | None = None, img_window_ratio: int = 32) → None

Method updates the image resolution to be processed and window size and so the pair-wise relative positions.

Parameters:

• new_window_size (Optional[int]) – New window size, if None based on new_img_size // window_div

• new_img_size (Optional[Tuple[int, int]]) – New input resolution, if None current resolution is used

• img_window_ratio (int) – divisor for calculating window size from image size

group_matcher(coarse=False)

set_grad_checkpointing(enable=True)

get_classifier() → torch.nn.Module

Method returns the classification head of the model. :returns: Current classification head :rtype: head (nn.Module)

reset_classifier(num_classes: int, global_pool: str | None = None) → None

Method results the classification head

Parameters:

• num_classes (int) – Number of classes to be predicted

• global_pool (str) – Unused

class credit.models.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.DebuggerModel(image_height: int = 640, image_width: int = 1280, frames: int = 2, channels: int = 4, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, levels: int = 15, post_conf: dict = None, **kwargs)

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.

image_height = 640

image_width = 1280

frames = 2

channels = 4

surface_channels = 7

levels = 15

input_only_channels = 3

output_only_channels = 0

linear

use_post_block

forward(x)

forward that multiplies self.coef to the input used to test postblock and other model parts

class credit.models.CrossFormerWithNoise(noise_latent_dim=128, encoder_noise_factor=0.05, decoder_noise_factor=0.275, encoder_noise=True, freeze=True, **kwargs)

Bases: credit.models.crossformer.CrossFormer

CrossFormer variant with pixel-wise noise injection in both encoder and decoder stages.

noise_latent_dim

Dimensionality of the noise vector.

Type:

int

encoder_noise_factor

Initial scaling factor for encoder noise injection.

Type:

float

decoder_noise_factor

Initial scaling factor for decoder noise injection.

Type:

float

encoder_noise

Whether to apply noise injection in the encoder.

Type:

bool

freeze

Whether to freeze pre-trained model weights.

Type:

bool

noise_latent_dim = 128

encoder_noise = True

noise_inject1

noise_inject2

noise_inject3

forward(x, noise=None, forecast_step=None)

Forward pass through the CrossFormer with noise injection.

Parameters:

• x (Tensor) – Input tensor of shape (batch_size, channels, height, width).

• noise (Tensor, optional) – External noise tensor. If None, noise is sampled internally. Defaults to None.

Returns:

Output tensor after passing through the model.

Return type:

Tensor

class credit.models.CrossFormerDiffusion(self_condition: bool = False, condition: bool = True, image_height: int = 640, patch_height: int = 1, image_width: int = 1280, patch_width: int = 1, frames: int = 2, channels: int = 4, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, levels: int = 15, dim: tuple = (64, 128, 256, 512), depth: tuple = (2, 2, 8, 2), dim_head: int = 32, global_window_size: tuple = (5, 5, 2, 1), local_window_size: int = 10, cross_embed_kernel_sizes: tuple = ((4, 8, 16, 32), (2, 4), (2, 4), (2, 4)), cross_embed_strides: tuple = (4, 2, 2, 2), attn_dropout: float = 0.0, ff_dropout: float = 0.0, use_spectral_norm: bool = True, interp: bool = True, padding_conf: dict = None, post_conf: dict = None, **kwargs)

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.

image_height = 640

image_width = 1280

patch_height = 1

patch_width = 1

frames = 2

channels = 4

surface_channels = 7

levels = 15

use_spectral_norm = True

use_interp = True

use_padding

use_post_block

input_only_channels = 3

input_channels = 70

output_channels = 67

total_input_channels = 70

condition = True

self_condition = False

layers

cube_embedding

time_to_emb

time_emb_proj

up_block1

up_block2

up_block3

up_block4

forward(x, t, x_self_cond=False, x_cond=None)

class credit.models.UnetDiffusion(image_height: int = 640, image_width: int = 1280, init_dim=None, frames: int = 2, channels: int = 4, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, levels: int = 15, dim: tuple = (64, 128, 256, 512), depth: tuple = (2, 2, 8, 2), dim_head: int = 32, padding_conf: dict = None, post_conf: dict = None, dim_mults: tuple = (1, 2, 4, 8), conditional_dimensions: int = 0, learned_variance: bool = False, learned_sinusoidal_cond: bool = False, random_fourier_features: bool = False, learned_sinusoidal_dim: int = 16, sinusoidal_pos_emb_theta: int = 10000, dropout: float = 0.0, attn_dim_head: int = 32, attn_heads: int = 4, full_attn: dict = None, flash_attn: bool = False, self_condition: bool = False, condition: bool = False, *args, **kwargs)

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.

output_channels = 67

channels = 4

channels_out = 67

pre_out_dim = 67

self_condition = False

conditional_dimensions = 0

condition = False

image_height = 640

image_width = 1280

frames = 2

surface_channels = 7

levels = 15

use_post_block

use_padding

input_only_channels = 3

input_channels = 70

random_or_learned_sinusoidal_cond = False

time_mlp

downs

ups

mid_block1

mid_attn

mid_block2

final_res_block

final_conv

property downsample_factor

forward(x, time, x_self_cond=None, x_cond=None)

class credit.models.ModifiedGaussianDiffusion(*args, **kwargs)

Bases: GaussianDiffusion

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.

channels

history_len

criterion = None

load_loss(criterion)

forward(img, x_cond=None, *args, **kwargs)

p_losses(x_start, t, x_cond, noise=None, offset_noise_strength=None)

model_predictions(x, t, x_self_cond=None, x_cond=None, clip_x_start=False, rederive_pred_noise=False)

credit.models.logger

credit.models.model_types

credit.models.load_fsdp_or_checkpoint_policy(conf)

credit.models.load_model(conf, load_weights=False, model_name=False)

credit.models.load_model_name(conf, model_name, load_weights=False)