File size: 1,192 Bytes
439c523 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """Factory and checkpoint helpers for the project-local NowcastNet model."""
from __future__ import annotations
from types import SimpleNamespace
from typing import Mapping
import torch
from .nowcastnet import Net
def build_model(config, device: torch.device | str = "cpu") -> Net:
"""Build a model from a mapping or namespace without requiring OneScience."""
if isinstance(config, Mapping):
config = SimpleNamespace(**config)
config.device = torch.device(device)
config.evo_ic = config.total_length - config.input_length
config.gen_oc = config.total_length - config.input_length
config.ic_feature = config.ngf * 10
return Net(config).to(config.device)
def load_checkpoint(model: torch.nn.Module, path: str, device: torch.device | str = "cpu") -> Mapping:
try:
state = torch.load(path, map_location=device, weights_only=True)
except TypeError: # torch < 2.0
state = torch.load(path, map_location=device)
checkpoint = state if isinstance(state, Mapping) else {}
model_state = checkpoint["state_dict"] if "state_dict" in checkpoint else state
model.load_state_dict(model_state, strict=True)
return checkpoint
|