"""Utilities for computing MEPS datastore standardization statistics."""
# Standard library
import os
import subprocess
from argparse import ArgumentParser
from datetime import timedelta
from pathlib import Path
from typing import Any, cast
# Third-party
import torch
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm
# First-party
from neural_lam import WeatherDataset
from neural_lam.datastore import init_datastore
from neural_lam.utils import get_integer_time
[docs]
class PaddedWeatherDataset(torch.utils.data.Dataset):
"""Wrap :class:`WeatherDataset` to pad samples for distributed runners."""
def __init__(
self,
base_dataset: WeatherDataset,
world_size: int,
batch_size: int,
) -> None:
"""
Parameters
----------
base_dataset : WeatherDataset
Dataset to pad.
world_size : int
Total number of distributed ranks participating.
batch_size : int
Per-rank batch size.
"""
super().__init__()
self.base_dataset = base_dataset
self.world_size = world_size
self.batch_size = batch_size
self.total_samples = len(base_dataset)
self.padded_samples = (
(self.world_size * self.batch_size) - self.total_samples
) % self.world_size
def __getitem__( # ty: ignore[invalid-method-override]
self, idx: int
) -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]:
"""Return an item, repeating the final sample for padded indices.
Parameters
----------
idx : int
Index into the padded dataset.
Returns
-------
tuple of torch.Tensor
The sample, as returned by `WeatherDataset.__getitem__`.
"""
return self.base_dataset[min(idx, self.total_samples - 1)]
def __len__(self) -> int:
"""Return the padded dataset length."""
return self.total_samples + self.padded_samples
[docs]
def real_sample_masks(
self, sampler: DistributedSampler
) -> list[torch.Tensor]:
"""
Return one boolean mask per minibatch, ``True`` for non-padded rows.
``DistributedSampler`` stripes indices across ranks, so the padded
tail of this dataset lands on several ranks and cannot be sliced off
the rank-major gathered tensor. Each rank drops its padded rows with
these masks instead; the sampler order chunked by ``batch_size`` is
the order a ``DataLoader(drop_last=False)`` yields batches in.
Parameters
----------
sampler : DistributedSampler
Sampler of this rank's ``DataLoader`` over this dataset.
Returns
-------
list[torch.Tensor]
Boolean mask per minibatch.
"""
is_real = torch.tensor(list(sampler)) < self.total_samples
return list(is_real.split(self.batch_size))
[docs]
def get_rank() -> int:
"""
Return the rank inferred from SLURM or default to 0.
Returns
-------
int
The current process rank.
"""
return int(os.environ.get("SLURM_PROCID", 0))
[docs]
def get_world_size() -> int:
"""
Return the world size inferred from SLURM or default to 1.
Returns
-------
int
The total number of processes in the distributed group.
"""
return int(os.environ.get("SLURM_NTASKS", 1))
[docs]
def setup(rank: int, world_size: int) -> None: # pylint: disable=redefined-outer-name
"""
Initialize the distributed group.
Parameters
----------
rank : int
The rank of the current process.
world_size : int
The total number of processes.
Raises
------
RuntimeError
If ``SLURM_JOB_NODELIST`` is set but no hostnames can be retrieved.
"""
if "SLURM_JOB_NODELIST" in os.environ:
nodelist = os.environ["SLURM_JOB_NODELIST"]
hostnames = subprocess.check_output(
["scontrol", "show", "hostnames", nodelist],
)
hostname_lines = hostnames.decode("utf-8").splitlines()
if not hostname_lines:
raise RuntimeError(
f"SLURM_JOB_NODELIST is set to {repr(nodelist)}, but "
"'scontrol show hostnames' returned no hostnames. "
"Please check your SLURM job configuration."
)
master_node = hostname_lines[0].strip()
else:
print(
"\033[91mCareful, you are running this script with --distributed "
"without any scheduler. In most cases this will result in slower "
"execution and the --distributed flag should be removed.\033[0m"
)
master_node = "localhost"
os.environ["MASTER_ADDR"] = master_node
os.environ["MASTER_PORT"] = "12355"
dist.init_process_group(
"nccl" if torch.cuda.is_available() else "gloo",
rank=rank,
world_size=world_size,
)
if rank == 0:
print(
f"Initialized {dist.get_backend()} "
f"process group with world size {world_size}."
)
[docs]
def save_stats(
static_dir_path: str | Path,
means: list[torch.Tensor],
squares: list[torch.Tensor],
flux_means: list[torch.Tensor],
flux_squares: list[torch.Tensor],
filename_prefix: str,
) -> None:
"""
Aggregate running statistics and persist them to ``static_dir_path``.
Parameters
----------
static_dir_path : str or pathlib.Path
Directory where ``*.pt`` files should be written.
means : Sequence[torch.Tensor]
Shape ``(B, num_state_vars)``. Batch-wise means. Sequence of one or
more tensors each of shape ``(B_i, num_state_vars)``; concatenated /
stacked along the batch dim before reduction. Typical callers pass a
single already-gathered tensor. The list can have length 0 to skip
saving.
squares : Sequence[torch.Tensor]
Shape ``(B, num_state_vars)``. Batch-wise second moments. Sequence of
one or more tensors each of shape ``(B_i, num_state_vars)``;
concatenated / stacked along the batch dim before reduction. Typical
callers pass a single already-gathered tensor. The list can have
length 0 to skip saving.
flux_means : Sequence[torch.Tensor]
Shape ``(B,)``. Flux means. Sequence of one or more tensors each of
shape ``(B_i,)``; concatenated / stacked along the batch dim before
reduction. Typical callers pass a single already-gathered tensor.
The list can have length 0 to skip saving.
flux_squares : Sequence[torch.Tensor]
Shape ``(B,)``. Flux second moments. Sequence of one or more tensors
each of shape ``(B_i,)``; concatenated / stacked along the batch dim
before reduction. Typical callers pass a single already-gathered
tensor. The list can have length 0 to skip saving.
filename_prefix : str
Prefix (e.g., ``"parameter"`` or ``"diff"``) for saved tensors.
"""
means_tensor = (
torch.stack(means) if len(means) > 1 else means[0]
) # (B, num_state_vars,)
squares_tensor = (
torch.stack(squares) if len(squares) > 1 else squares[0]
) # (B, num_state_vars,)
mean = torch.mean(means_tensor, dim=0) # (num_state_vars,)
second_moment = torch.mean(squares_tensor, dim=0) # (num_state_vars,)
std = torch.sqrt(second_moment - mean**2) # (num_state_vars,)
print(
f"Saving {filename_prefix} mean and std.-dev. to "
f"{filename_prefix}_mean.pt and {filename_prefix}_std.pt"
)
torch.save(
mean.cpu(), os.path.join(static_dir_path, f"{filename_prefix}_mean.pt")
)
torch.save(
std.cpu(), os.path.join(static_dir_path, f"{filename_prefix}_std.pt")
)
if len(flux_means) == 0:
return
flux_means_tensor = (
torch.stack(flux_means) if len(flux_means) > 1 else flux_means[0]
) # (B,)
flux_squares_tensor = (
torch.stack(flux_squares) if len(flux_squares) > 1 else flux_squares[0]
) # (B,)
flux_mean = torch.mean(flux_means_tensor) # (,)
flux_second_moment = torch.mean(flux_squares_tensor) # (,)
flux_std = torch.sqrt(flux_second_moment - flux_mean**2) # (,)
print("Saving flux mean and std.-dev. to flux_stats.pt")
torch.save(
torch.stack((flux_mean, flux_std)).cpu(),
os.path.join(static_dir_path, "flux_stats.pt"),
)
[docs]
def main(
datastore_config_path: str | Path,
batch_size: int,
step_length: timedelta,
n_workers: int,
distributed: bool,
) -> None:
"""
Pre-compute and persist standardization statistics from the datastore.
Parameters
----------
datastore_config_path : str or pathlib.Path
Path to the MEPS datastore configuration file.
batch_size : int
Batch size used while iterating through the dataset.
step_length : datetime.timedelta
Temporal sampling interval for the difference statistics.
n_workers : int
Number of dataloader workers.
distributed : bool
If ``True``, run using torch.distributed with SLURM settings.
"""
rank = get_rank()
world_size = get_world_size()
datastore = init_datastore(
datastore_kind="npyfilesmeps", config_path=datastore_config_path
)
static_dir_path = Path(datastore_config_path).parent / "static"
os.makedirs(static_dir_path, exist_ok=True)
if distributed:
setup(rank, world_size)
device = torch.device(
f"cuda:{rank}" if torch.cuda.is_available() else "cpu"
)
torch.cuda.set_device(device) if torch.cuda.is_available() else None
# Setting this to the original value of the Oskarsson et al. paper (2023)
# 65 forecast steps - 2 initial steps = 63
ar_steps = 63
# Raw (non-standardized) data for computing mean/std
ds_base = WeatherDataset(
datastore=datastore,
split="train",
ar_steps=ar_steps,
num_past_forcing_steps=0,
num_future_forcing_steps=0,
)
sampler: DistributedSampler | None
ds: torch.utils.data.Dataset
if distributed:
ds = PaddedWeatherDataset(
ds_base,
world_size,
batch_size,
)
sampler = DistributedSampler(
ds, num_replicas=world_size, rank=rank, shuffle=False
)
else:
ds = ds_base
sampler = None
loader = torch.utils.data.DataLoader(
ds,
batch_size,
shuffle=False,
num_workers=n_workers,
sampler=sampler,
)
if rank == 0:
print("Computing mean and std.-dev. for parameters...")
means, squares, flux_means, flux_squares = [], [], [], []
real_masks = None
if distributed:
assert isinstance(ds, PaddedWeatherDataset)
assert sampler is not None
real_masks = ds.real_sample_masks(sampler)
for batch_i, (init_batch, target_batch, forcing_batch, _, _) in enumerate(
tqdm(loader)
):
if distributed:
init_batch, target_batch, forcing_batch = (
init_batch.to(device),
target_batch.to(device),
forcing_batch.to(device),
)
# (B, num_times, num_grid_nodes, num_state_vars)
batch = torch.cat((init_batch, target_batch), dim=1)
# Flux at 1st windowed position is index 0 in forcing
flux_batch = forcing_batch[:, :, :, 0]
# (B, num_state_vars,)
batch_means = torch.mean(batch, dim=(1, 2)).cpu()
batch_squares = torch.mean(batch**2, dim=(1, 2)).cpu()
# (B,)
batch_flux_means = torch.mean(flux_batch, dim=(1, 2)).cpu()
batch_flux_squares = torch.mean(flux_batch**2, dim=(1, 2)).cpu()
if real_masks is not None:
mask = real_masks[batch_i]
batch_means = batch_means[mask]
batch_squares = batch_squares[mask]
batch_flux_means = batch_flux_means[mask]
batch_flux_squares = batch_flux_squares[mask]
means.append(batch_means)
squares.append(batch_squares)
flux_means.append(batch_flux_means)
flux_squares.append(batch_flux_squares)
if distributed and world_size > 1:
means_gathered: list[Any] = [None] * world_size
squares_gathered: list[Any] = [None] * world_size
flux_means_gathered: list[Any] = [None] * world_size
flux_squares_gathered: list[Any] = [None] * world_size
dist.all_gather_object(means_gathered, torch.cat(means, dim=0))
dist.all_gather_object(squares_gathered, torch.cat(squares, dim=0))
dist.all_gather_object(
flux_means_gathered, torch.cat(flux_means, dim=0)
)
dist.all_gather_object(
flux_squares_gathered, torch.cat(flux_squares, dim=0)
)
if rank == 0:
means = [
torch.cat(cast("list[torch.Tensor]", means_gathered), dim=0)
]
squares = [
torch.cat(cast("list[torch.Tensor]", squares_gathered), dim=0)
]
flux_means = [
torch.cat(
cast("list[torch.Tensor]", flux_means_gathered), dim=0
)
]
flux_squares = [
torch.cat(
cast("list[torch.Tensor]", flux_squares_gathered), dim=0
)
]
else:
means = [torch.cat(means, dim=0)] # (B, num_state_vars,)
squares = [torch.cat(squares, dim=0)] # (B, num_state_vars,)
flux_means = [torch.cat(flux_means, dim=0)] # (B,)
flux_squares = [torch.cat(flux_squares, dim=0)] # (B,)
if rank == 0:
save_stats(
static_dir_path,
means,
squares,
flux_means,
flux_squares,
"parameter",
)
if distributed:
dist.barrier()
if rank == 0:
print("Computing mean and std.-dev. for one-step differences...")
ds_standard_base = WeatherDataset(
datastore=datastore,
split="train",
ar_steps=ar_steps,
num_past_forcing_steps=0,
num_future_forcing_steps=0,
)
sampler_standard: DistributedSampler | None
ds_standard: torch.utils.data.Dataset
if distributed:
ds_standard = PaddedWeatherDataset(
ds_standard_base,
world_size,
batch_size,
)
sampler_standard = DistributedSampler(
ds_standard, num_replicas=world_size, rank=rank, shuffle=False
)
else:
ds_standard = ds_standard_base
sampler_standard = None
loader_standard = torch.utils.data.DataLoader(
ds_standard,
batch_size,
shuffle=False,
num_workers=n_workers,
sampler=sampler_standard,
)
# WeatherDataset no longer standardizes, so load the state mean/std saved
# in the first pass above and apply them inline to compute diff stats.
state_mean = torch.load(
os.path.join(static_dir_path, "parameter_mean.pt"), weights_only=True
)
state_std = torch.load(
os.path.join(static_dir_path, "parameter_std.pt"), weights_only=True
)
if distributed:
state_mean = state_mean.to(device)
state_std = state_std.to(device)
time_step_int, time_step_unit = get_integer_time(step_length)
assert time_step_unit == "hours", (
"Only 'hours' time unit is supported by meps datastore."
)
used_subsample_len = (65 // time_step_int) * time_step_int
diff_means, diff_squares = [], []
real_masks_standard = None
if distributed:
assert isinstance(ds_standard, PaddedWeatherDataset)
assert sampler_standard is not None
real_masks_standard = ds_standard.real_sample_masks(sampler_standard)
for batch_i, (init_batch, target_batch, _, _, _) in enumerate(
tqdm(loader_standard, disable=rank != 0)
):
if distributed:
init_batch, target_batch = (
init_batch.to(device),
target_batch.to(device),
)
init_batch = (init_batch - state_mean) / state_std
target_batch = (target_batch - state_mean) / state_std
# (B, num_times', num_grid_nodes, num_state_vars)
batch = torch.cat((init_batch, target_batch), dim=1)
# Note: batch contains only 1h-steps
stepped_batch = torch.cat(
[
batch[:, ss_i:used_subsample_len:time_step_int]
for ss_i in range(time_step_int)
],
dim=0,
)
# (B', num_times, num_grid_nodes, num_state_vars),
# B' = step_length*B
batch_diffs = stepped_batch[:, 1:] - stepped_batch[:, :-1]
# (B', num_times-1, num_grid_nodes, num_state_vars)
batch_diff_means = torch.mean(batch_diffs, dim=(1, 2)).cpu()
# (B', num_state_vars,)
batch_diff_squares = torch.mean(batch_diffs**2, dim=(1, 2)).cpu()
# (B', num_state_vars,)
if real_masks_standard is not None:
# stepped_batch stacks the time_step_int ss_i blocks along dim 0,
# so tile the mask (repeat, not repeat_interleave)
mask = real_masks_standard[batch_i].repeat(time_step_int)
batch_diff_means = batch_diff_means[mask]
batch_diff_squares = batch_diff_squares[mask]
diff_means.append(batch_diff_means)
diff_squares.append(batch_diff_squares)
if distributed and world_size > 1:
dist.barrier()
diff_means_gathered: list[Any] = [None] * world_size
diff_squares_gathered: list[Any] = [None] * world_size
dist.all_gather_object(
diff_means_gathered, torch.cat(diff_means, dim=0)
)
dist.all_gather_object(
diff_squares_gathered, torch.cat(diff_squares, dim=0)
)
if rank == 0:
diff_means = [
torch.cat(
cast("list[torch.Tensor]", diff_means_gathered), dim=0
)
]
diff_squares = [
torch.cat(
cast("list[torch.Tensor]", diff_squares_gathered), dim=0
)
]
diff_means = [torch.cat(diff_means, dim=0)] # (B', num_state_vars,)
diff_squares = [torch.cat(diff_squares, dim=0)] # (B', num_state_vars,)
if rank == 0:
save_stats(static_dir_path, diff_means, diff_squares, [], [], "diff")
if distributed:
dist.destroy_process_group()
[docs]
def cli() -> None:
"""Parse CLI arguments and trigger :func:`main`."""
parser = ArgumentParser(description="Training arguments")
parser.add_argument(
"--datastore_config_path",
type=str,
help="Path to data config file",
)
parser.add_argument(
"--batch_size",
type=int,
default=32,
help="Batch size when iterating over the dataset",
)
parser.add_argument(
"--step_length",
type=int,
default=3,
help="Step length in hours to consider single time step (default: 3)",
)
parser.add_argument(
"--n_workers",
type=int,
default=4,
help="Number of workers in data loader (default: 4)",
)
parser.add_argument(
"--distributed",
action="store_true",
help="Run the script in distributed mode (default: False)",
)
args = parser.parse_args()
distributed = bool(args.distributed)
main(
datastore_config_path=args.datastore_config_path,
batch_size=args.batch_size,
step_length=timedelta(hours=args.step_length),
n_workers=args.n_workers,
distributed=distributed,
)
if __name__ == "__main__":
cli()
