Skip to content

Base Model

ocean_emulators.models.base

BaseModel(in_channels, out_channels, hist, pred_residuals, last_kernel_size, pad, gradient_detach_interval)

Bases: Module

Abstract base model for neural ocean emulators.

Provides shared functionality for all ocean emulator models, including residual prediction, ocean masking, and gradient detaching for multi-step autoregressive rollouts.

Source code in src/ocean_emulators/models/base.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(
    self,
    in_channels,
    out_channels,
    hist,
    pred_residuals,
    last_kernel_size,
    pad,
    gradient_detach_interval: int,
) -> None:
    super().__init__()
    assert last_kernel_size % 2 != 0, "Cannot use even kernel sizes!"
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.N_pad = int((last_kernel_size - 1) / 2)
    self.pad: str = pad
    self.pred_residuals = pred_residuals
    self.hist = hist
    self.gradient_detach_interval = gradient_detach_interval