library_name: transformers
tags:
- weathernext2
- tiny-random
- testing
Tiny random WeatherNext 2
A randomly initialized WeatherNext2ForWeatherForecasting
for use in the Transformers test suite. 13K parameters, 152 KB. It forecasts nothing: the weights are
noise and the feature extractor's statistics are zero mean and unit variance.
What it is for is the geometry. WeatherNext 2 has no learned positional encodings, so the mesh,
both grid-mesh graphs and the banded attention mask are built once at conversion time, with scipy
and trimesh, and stored in the checkpoint as buffers. A model built from a config alone gets a
placeholder instead, which is enough to run but is not a mesh. This checkpoint carries the real
thing at a size CI can download, so the tests that depend on a real mesh have one.
The grid is 15 degrees and the mesh is a twice-split icosahedron: 162 mesh nodes over 312 grid points, an attention bandwidth of 43, so four attention blocks. Four is the point, since it gives an interior block with a neighbour on either side.
Regenerating it
import torch
from transformers import WeatherNext2Config, WeatherNext2FeatureExtractor, WeatherNext2ForWeatherForecasting
from transformers.models.weathernext2.convert_weathernext2_original_checkpoint import geometry_state_dict
config = WeatherNext2Config(
hidden_size=16, intermediate_size=32, num_hidden_layers=2, num_attention_heads=2,
edge_hidden_size=4, noise_channels=4, mesh_splits=2, attention_k_hop=2,
grid_latitudes=13, grid_longitudes=24, pressure_levels=(500, 850), aggregate_normalization=None,
)
geometry = geometry_state_dict(config) # needs scipy and trimesh; sets the two derived sizes
torch.manual_seed(0)
model = WeatherNext2ForWeatherForecasting(config)
model.load_state_dict(geometry, strict=False)
The feature extractor is built from the same config with zero means and unit standard deviations for every variable.