dynabench.model

Module containing the models used in the DynaBench benchmark.

Classes

CNN(input_size, output_size[, ...])

Simple 2D CNN model for grid data.

NeuralPDE(input_dim[, hidden_channels, ...])

Neural PDE model for grid data.

ResNet(input_size, output_size[, ...])

Simple 2D ResNet model for grid data.
class dynabench.model.CNN(input_size: int, output_size: int, hidden_layers: int = 1, hidden_channels: int = 64, padding: int | str | Tuple[int] = 'same', padding_mode: str = 'circular', kernel_size: int = 3, activation: str = 'ReLU')[source]

Bases: Module

Simple 2D CNN model for grid data.

Parameters:

  • input_size (int) – Number of input channels.

  • output_size (int) – Number of output channels.

  • hidden_layers (int) – Number of hidden layers. Default is 1.

  • hidden_channels (int) – Number of channels in each hidden layer. Default is 64.

  • padding (int | str | Tuple[int]) – Padding size. If ‘same’, padding is calculated to keep the input size the same as the output size. Default is ‘same’.

  • padding_mode (str) – What value to pad with. Can be ‘zeros’, ‘reflect’, ‘replicate’ or ‘circular’. Default: ‘zeros’

  • kernel_size (int) – Size of the kernel. Default is 3.

  • activation (str) – Activation function to use. Can be one of torch.nn activation functions. Default is ‘relu’.

forward(x: Tensor)[source]

Forward pass of the model. Should not be called directly, instead call the model instance.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, input_size, height, width).

Returns:

Output tensor of shape (batch_size, output_size, height, width).

Return type:

torch.Tensor

class dynabench.model.NeuralPDE(input_dim: int, hidden_channels: int = 64, hidden_layers: int = 1, solver: dict = {'method': 'dopri5'}, use_adjoint: bool = True, *args, **kwargs)[source]

Bases: Module

Neural PDE model for grid data. The model combines a CNN with a differentiable ODE solver to learn the dynamics of the data using the method of lines. The CNN is used to approximate the spatial derivatives of the data, while the ODE solver is used to approximate the temporal evolution of the data. The model has been taken from NeuralPDE: Modelling Dynamical Systems from Data by Dulny et al.

Parameters:

  • input_dim (int) – Number of input channels.

  • hidden_channels (int) – Number of channels in each hidden layer of the CNN. Default is 64.

  • hidden_layers (int) – Number of hidden layers in the CNN. Default is 1.

  • solver (dict) – Dictionary of solver parameters. Default is {“method”: “dopri5”}.

  • use_adjoint (bool) – Whether to use the adjoint method for backpropagation. Default is True.

forward(x: Tensor, t_eval: List[float] = [0.0, 1.0])[source]

Forward pass of the model. Should not be called directly, instead call the model instance.

Parameters:

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

  • t_eval (List[float], default [0.0, 1.0]) – List of times to evaluate the ODE solver at. Default is [0.0, 1.0].

Returns:

Output tensor of shape (batch_size, rollout, output_size, height, width).

Return type:

torch.Tensor

class dynabench.model.ResNet(input_size: int, output_size: int, resblock_layers: int = 1, resblock_channels: int = 64, resblock_type: str = 'simple', padding: int | str | Tuple[int] = 'same', padding_mode: str = 'circular', kernel_size: int = 3, activation: str = 'ReLU')[source]

Bases: Module

Simple 2D ResNet model for grid data.

Parameters:

  • input_size (int) – Number of input channels.

  • output_size (int) – Number of output channels.

  • resblock_layers (int) – Number of residual blocks. Default is 1.

  • resblock_channels (int) – Number of channels in each residual block. Default is 64.

  • resblock_type (str) – Type of residual block to use. Can be ‘simple’ or ‘bottleneck’. Default is ‘simple’.

  • padding (int | str | Tuple[int]) – Padding size. If ‘same’, padding is calculated to keep the input size the same as the output size. Default is ‘same’.

  • padding_mode (str) – What value to pad with. Can be ‘zeros’, ‘reflect’, ‘replicate’ or ‘circular’. Default: ‘zeros’

  • kernel_size (int) – Size of the kernel. Default is 3.

  • activation (str) –

    Activation function to use. Can be one of torch.nn activation functions. Default is ‘relu’.

forward(x)[source]

Forward pass of the model. Should not be called directly, instead call the model instance.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, input_size, height, width).

Returns:

Output tensor of shape (batch_size, output_size, height, width).

Return type:

torch.Tensor

Modules

dynabench.model.grind