src.FRAME_FM.models.demo_simple_conv

Attributes

KERCHUNK_ZIP
n_variables

Classes

MyModel	A thin wrapper around PyTorch Lightning's LightningModule to allow for future extensions
SpatialCollapse	Input: [B, V, H, W]

Module Contents

class src.FRAME_FM.models.demo_simple_conv.MyModel

Bases: FRAME_FM.utils.LightningModuleWrapper.BaseModule

A thin wrapper around PyTorch Lightning’s LightningModule to allow for future extensions and customizations specific to FRAME-FM project needs.

Subclasses should implement training_step_body and validation_step_body methods instead of training_step and validation_step directly.

Enforces consistent logging patterns across training and validation steps. Enforces logging of loss by default.

class src.FRAME_FM.models.demo_simple_conv.SpatialCollapse(in_variables, hidden_channels=32, dropout=0.2)

Bases: torch.nn.Module

Input: [B, V, H, W] Output: [B]

conv1

bn1

conv2

bn2

conv3

bn3

dropout

pool

head

forward(x)

src.FRAME_FM.models.demo_simple_conv.KERCHUNK_ZIP = 'tests/transforms/fixtures/ecmwf-era5X_oper_an_sfc_2000_2020_2d_repack.kr1.0.json.zip'

src.FRAME_FM.models.demo_simple_conv.n_variables = 5