credit.models ============= .. py:module:: credit.models Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/credit/models/base_model/index /autoapi/credit/models/checkpoint/index /autoapi/credit/models/crossformer/index /autoapi/credit/models/crossformer_diffusion/index /autoapi/credit/models/crossformer_ensemble/index /autoapi/credit/models/debugger_model/index /autoapi/credit/models/fuxi/index /autoapi/credit/models/graph/index /autoapi/credit/models/reset/index /autoapi/credit/models/swin/index /autoapi/credit/models/unet/index /autoapi/credit/models/unet404/index Attributes ---------- .. autoapisummary:: credit.models.logger credit.models.model_types Classes ------- .. autoapisummary:: credit.models.CrossFormer credit.models.SegmentationModel credit.models.SegmentationModel404 credit.models.Fuxi credit.models.SwinTransformerV2Cr credit.models.GraphResTransfGRU credit.models.DebuggerModel credit.models.CrossFormerWithNoise Functions --------- .. autoapisummary:: credit.models.load_fsdp_or_checkpoint_policy credit.models.load_model credit.models.load_model_name Package Contents ---------------- .. py:class:: 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, interp: bool = True, padding_conf: dict = None, post_conf: dict = None, **kwargs) 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:: image_height :value: 640 .. py:attribute:: image_width :value: 1280 .. py:attribute:: patch_height :value: 1 .. py:attribute:: patch_width :value: 1 .. py:attribute:: frames :value: 2 .. py:attribute:: channels :value: 4 .. py:attribute:: surface_channels :value: 7 .. py:attribute:: levels :value: 15 .. py:attribute:: use_spectral_norm :value: True .. py:attribute:: use_interp :value: True .. py:attribute:: use_padding .. py:attribute:: use_post_block .. py:attribute:: input_only_channels :value: 3 .. py:attribute:: input_channels :value: 70 .. py:attribute:: output_channels :value: 67 .. py:attribute:: layers .. py:attribute:: cube_embedding .. py:attribute:: up_block1 .. py:attribute:: up_block2 .. py:attribute:: up_block3 .. py:attribute:: up_block4 .. py:method:: forward(x) .. py:method:: rk4(x) .. py:class:: 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: :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:: image_height :value: 640 .. py:attribute:: image_width :value: 1280 .. py:attribute:: frames :value: 2 .. py:attribute:: channels :value: 4 .. py:attribute:: surface_channels :value: 7 .. py:attribute:: levels :value: 16 .. py:attribute:: rk4_integration :value: False .. py:attribute:: model .. py:attribute:: use_post_block .. py:method:: forward(x) .. py:method:: rk4(x) .. py:class:: SegmentationModel404(conf) 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:: variables .. py:attribute:: frames .. py:attribute:: static_variables .. py:attribute:: model .. py:method:: forward(x) .. py:class:: 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: :py:obj:`credit.models.base_model.BaseModel` :param img_size: T, Lat, Lon. :type img_size: Sequence[int], optional :param patch_size: T, Lat, Lon. :type patch_size: Sequence[int], optional :param in_chans: number of input channels. :type in_chans: int, optional :param out_chans: number of output channels. :type out_chans: int, optional :param dim: number of embed channels. :type dim: int, optional :param num_groups: number of groups to separate the channels into. :type num_groups: Sequence[int] | int, optional :param num_heads: Number of attention heads. :type num_heads: int, optional :param window_size: Local window size. :type window_size: int | tuple[int], optional .. py:attribute:: use_interp :value: True .. py:attribute:: use_spectral_norm :value: True .. py:attribute:: use_padding .. py:attribute:: use_post_block .. py:attribute:: cube_embedding .. py:attribute:: u_transformer .. py:attribute:: fc .. py:attribute:: patch_size .. py:attribute:: input_resolution .. py:attribute:: out_chans :value: 67 .. py:attribute:: img_size .. py:attribute:: channels :value: 4 .. py:attribute:: surface_channels :value: 7 .. py:attribute:: levels :value: 15 .. py:method:: forward(x: torch.Tensor) .. py:class:: SwinTransformerV2Cr(img_size: Tuple[int, int] = (224, 224), patch_size: int = 4, window_size: Optional[int] = 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: Optional[float] = 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: :py:obj:`credit.models.base_model.BaseModel` Swin Transformer V2 A PyTorch impl of : `Swin Transformer V2: Scaling Up Capacity and Resolution` - https://arxiv.org/pdf/2111.09883 :param img_size: Input resolution. :param window_size: Window size. If None, img_size // window_div :param img_window_ratio: Window size to image size ratio. :param patch_size: Patch size. :param in_chans: Number of input channels. :param depths: Depth of the stage (number of layers). :param num_heads: Number of attention heads to be utilized. :param embed_dim: Patch embedding dimension. :param num_classes: Number of output classes. :param mlp_ratio: Ratio of the hidden dimension in the FFN to the input channels. :param drop_rate: Dropout rate. :param proj_drop_rate: Projection dropout rate. :param attn_drop_rate: Dropout rate of attention map. :param drop_path_rate: Stochastic depth rate. :param norm_layer: Type of normalization layer to be utilized. :param extra_norm_period: Insert extra norm layer on main branch every N (period) blocks in stage :param extra_norm_stage: End each stage with an extra norm layer in main branch :param sequential_attn: If true sequential self-attention is performed. :param padding_conf: padding configuration :type padding_conf: dict :param post_conf: configuration for postblock processing :type post_conf: dict .. py:attribute:: use_padding .. py:attribute:: patch_size :type: int :value: 4 .. py:attribute:: img_size :type: Tuple[int, int] :value: (224, 224) .. py:attribute:: window_size :type: int .. py:attribute:: num_features :type: int :value: 96 .. py:attribute:: frames :value: 1 .. py:attribute:: in_chans :value: 70 .. py:attribute:: out_chans :value: 67 .. py:attribute:: feature_info :value: [] .. py:attribute:: full_pos_embed :value: False .. py:attribute:: checkpoint_stages :value: False .. py:attribute:: residual :value: False .. py:attribute:: depth .. py:attribute:: use_post_block .. py:attribute:: patch_embed .. py:attribute:: stages .. py:attribute:: head .. py:attribute:: use_spectral_norm :value: False .. py:method:: forward_features(x: torch.Tensor) -> torch.Tensor .. py:method:: forward_head(x: torch.Tensor) -> torch.Tensor .. py:method:: forward(x: torch.Tensor) -> torch.Tensor .. py:method:: update_input_size(new_img_size: Optional[Tuple[int, int]] = None, new_window_size: Optional[int] = 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. :param new_window_size: New window size, if None based on new_img_size // window_div :type new_window_size: Optional[int] :param new_img_size: New input resolution, if None current resolution is used :type new_img_size: Optional[Tuple[int, int]] :param img_window_ratio: divisor for calculating window size from image size :type img_window_ratio: int .. py:method:: group_matcher(coarse=False) .. py:method:: set_grad_checkpointing(enable=True) .. py:method:: get_classifier() -> torch.nn.Module Method returns the classification head of the model. :returns: Current classification head :rtype: head (nn.Module) .. py:method:: reset_classifier(num_classes: int, global_pool: Optional[str] = None) -> None Method results the classification head :param num_classes: Number of classes to be predicted :type num_classes: int :param global_pool: Unused :type global_pool: str .. 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:: 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: :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:: image_height :value: 640 .. py:attribute:: image_width :value: 1280 .. py:attribute:: frames :value: 2 .. py:attribute:: channels :value: 4 .. py:attribute:: surface_channels :value: 7 .. py:attribute:: levels :value: 15 .. py:attribute:: input_only_channels :value: 3 .. py:attribute:: output_only_channels :value: 0 .. py:attribute:: linear .. py:attribute:: use_post_block .. py:method:: forward(x) forward that multiplies self.coef to the input used to test postblock and other model parts .. py:class:: CrossFormerWithNoise(noise_latent_dim=128, encoder_noise_factor=0.05, decoder_noise_factor=0.275, encoder_noise=True, freeze=True, **kwargs) Bases: :py:obj:`credit.models.crossformer.CrossFormer` CrossFormer variant with pixel-wise noise injection in both encoder and decoder stages. .. attribute:: noise_latent_dim Dimensionality of the noise vector. :type: int .. attribute:: encoder_noise_factor Initial scaling factor for encoder noise injection. :type: float .. attribute:: decoder_noise_factor Initial scaling factor for decoder noise injection. :type: float .. attribute:: encoder_noise Whether to apply noise injection in the encoder. :type: bool .. attribute:: freeze Whether to freeze pre-trained model weights. :type: bool .. py:attribute:: noise_latent_dim :value: 128 .. py:attribute:: encoder_noise :value: True .. py:attribute:: noise_inject1 .. py:attribute:: noise_inject2 .. py:attribute:: noise_inject3 .. py:method:: forward(x, noise=None, forecast_step=None) Forward pass through the CrossFormer with noise injection. :param x: Input tensor of shape (batch_size, channels, height, width). :type x: Tensor :param noise: External noise tensor. If None, noise is sampled internally. Defaults to None. :type noise: Tensor, optional :returns: Output tensor after passing through the model. :rtype: Tensor .. py:data:: logger .. py:data:: model_types .. py:function:: load_fsdp_or_checkpoint_policy(conf) .. py:function:: load_model(conf, load_weights=False, model_name=False) .. py:function:: load_model_name(conf, model_name, load_weights=False)