weather_projector

This model takes weather/2d inputs (batch, features, height, width) and produces tokens for multimodal language models.

class mfai.pytorch.models.weather_projector.PatchMaker(patch_size, input_dims, autopadding=True)[source]

Bases: Module

Converts a vision/weather (B, C, H, W) tensor into a (B, T, F) token tensor. T stands for token dimension and F the feature/embedding dimension. Each token is built with all the data of one patch of size patch_size.

Parameters:
forward(t)[source]
Return type:

Tensor

Parameters:

t (Tensor)

  1. zero pad if padding is enabled

  2. check for dim consistency

  3. einops rearrange to patch.

class mfai.pytorch.models.weather_projector.WeatherProjector(settings=WeatherProjectorSettings(patch_size=None, input_dims=(3, 256, 256), embedding_dim=768))[source]

Bases: Module

Parameters:

settings (WeatherProjectorSettings)

forward(t)[source]

Forward function of the WeatherProjector vision encoder.

Parameters:

t (Tensor) – tensor of shape (B, nu_channels, height, width)

Returns:

tensor of shape (B, num_patches_h * num_patches_w, embed_dim)

Return type:

Tensor

class mfai.pytorch.models.weather_projector.WeatherProjectorSettings(patch_size=None, input_dims=(3, 256, 256), embedding_dim=768)[source]

Bases: object

Parameters:
embedding_dim: int = 768
classmethod from_dict(kvs, *, infer_missing=False)
Return type:

TypeVar(A, bound= DataClassJsonMixin)

Parameters:

kvs (dict | list | str | int | float | bool | None)

classmethod from_json(s, *, parse_float=None, parse_int=None, parse_constant=None, infer_missing=False, **kw)
Return type:

TypeVar(A, bound= DataClassJsonMixin)

Parameters:

s (str | bytes | bytearray)

input_dims: tuple[int, int, int] = (3, 256, 256)
patch_size: None | int | tuple[int, int] = None
classmethod schema(*, infer_missing=False, only=None, exclude=(), many=False, context=None, load_only=(), dump_only=(), partial=False, unknown=None)
Return type:

SchemaF[TypeVar(A, bound= DataClassJsonMixin)]

Parameters:
to_dict(encode_json=False)
Return type:

Dict[str, Union[dict, list, str, int, float, bool, None]]

to_json(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, indent=None, separators=None, default=None, sort_keys=False, **kw)
Return type:

str

Parameters: