credit.models.fuxi ================== .. py:module:: credit.models.fuxi Attributes ---------- .. autoapisummary:: credit.models.fuxi.logger credit.models.fuxi.image_height Classes ------- .. autoapisummary:: credit.models.fuxi.CubeEmbedding credit.models.fuxi.DownBlock credit.models.fuxi.UpBlock credit.models.fuxi.UTransformer credit.models.fuxi.Fuxi Functions --------- .. autoapisummary:: credit.models.fuxi.apply_spectral_norm credit.models.fuxi.get_pad3d credit.models.fuxi.get_pad2d Module Contents --------------- .. py:data:: logger .. py:function:: apply_spectral_norm(model) add spectral norm to all the conv and linear layers .. py:function:: get_pad3d(input_resolution, window_size) Estimate the size of padding based on the given window size and the original input size. :param input_resolution: (Pl, Lat, Lon) :type input_resolution: tuple[int] :param window_size: (Pl, Lat, Lon) :type window_size: tuple[int] :returns: (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back) :rtype: padding (tuple[int]) .. py:function:: get_pad2d(input_resolution, window_size) :param input_resolution: Lat, Lon :type input_resolution: tuple[int] :param window_size: Lat, Lon :type window_size: tuple[int] :returns: (padding_left, padding_right, padding_top, padding_bottom) :rtype: padding (tuple[int]) .. py:class:: CubeEmbedding(img_size, patch_size, in_chans, embed_dim, norm_layer=nn.LayerNorm) Bases: :py:obj:`torch.nn.Module` :param img_size: T, Lat, Lon :param patch_size: T, Lat, Lon .. py:attribute:: img_size .. py:attribute:: patches_resolution .. py:attribute:: embed_dim .. py:attribute:: proj .. py:method:: forward(x: torch.Tensor) .. py:class:: DownBlock(in_chans: int, out_chans: int, num_groups: int, num_residuals: int = 2) 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:: conv .. py:attribute:: b .. py:method:: forward(x) .. py:class:: UpBlock(in_chans, out_chans, num_groups, num_residuals=2) 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:: conv .. py:attribute:: b .. py:method:: forward(x) .. py:class:: UTransformer(embed_dim, num_groups, input_resolution, num_heads, window_size, depth, proj_drop, attn_drop, drop_path, noise_latent_dim=128, noise_factor=0.2, use_noise=False, noise_scheduler=None) Bases: :py:obj:`torch.nn.Module` U-Transformer :param embed_dim: Patch embedding dimension. :type embed_dim: int :param num_groups: number of groups to separate the channels into. :type num_groups: int | tuple[int] :param input_resolution: Lat, Lon. :type input_resolution: tuple[int] :param num_heads: Number of attention heads in different layers. :type num_heads: int :param window_size: Window size. :type window_size: int | tuple[int] :param depth: Number of blocks. :type depth: int .. py:attribute:: padding .. py:attribute:: pad .. py:attribute:: down .. py:attribute:: layer .. py:attribute:: up .. py:attribute:: use_noise :value: False .. py:method:: forward(x, noise=None, forecast_step=None) .. 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, use_noise=False, noise_latent_dim=128, noise_factor=0.2, noise_scheduler=None, freeze=False, **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, noise=None, forecast_step=None) .. py:data:: image_height :value: 640