File size: 3,066 Bytes
8b16bae | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | from __future__ import annotations
import random
from collections.abc import Mapping
from pathlib import Path
from typing import Any
import numpy as np
import torch
import yaml
class AttrDict(dict):
"""Dictionary with recursive attribute access for OneScience datapipes."""
def __getattr__(self, key: str) -> Any:
try:
return self[key]
except KeyError as error:
raise AttributeError(key) from error
def __setattr__(self, key: str, value: Any) -> None:
self[key] = value
def to_attr_dict(value: Any) -> Any:
if isinstance(value, Mapping):
return AttrDict({key: to_attr_dict(item) for key, item in value.items()})
if isinstance(value, list):
return [to_attr_dict(item) for item in value]
return value
def to_plain_dict(value: Any) -> Any:
if isinstance(value, Mapping):
return {key: to_plain_dict(item) for key, item in value.items()}
if isinstance(value, list):
return [to_plain_dict(item) for item in value]
return value
def load_config(path: Path) -> dict[str, Any]:
with path.open("r", encoding="utf-8") as stream:
config = yaml.safe_load(stream)
if not isinstance(config, dict) or not isinstance(config.get("root"), dict):
raise ValueError(f"Config must contain a root mapping: {path}")
return config["root"]
def project_path(value: str | Path, project_root: Path) -> Path:
path = Path(value).expanduser()
return path if path.is_absolute() else project_root / path
def resolve_device(requested: str) -> torch.device:
if requested == "auto":
return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device(requested)
if device.type == "cuda" and not torch.cuda.is_available():
raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
return device
def seed_everything(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def rollout(
model: torch.nn.Module,
pos: torch.Tensor,
state: torch.Tensor,
t_out: int,
out_dim: int,
max_latent_steps: int,
) -> torch.Tensor:
predictions = []
remaining = int(t_out)
while remaining > 0:
step_count = min(int(max_latent_steps), remaining)
prediction = model(pos, state, latent_steps=step_count)
predictions.append(prediction)
consumed = step_count * int(out_dim)
state = torch.cat((state[..., consumed:], prediction), dim=-1)
remaining -= step_count
return torch.cat(predictions, dim=-1)
def relative_l2(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
prediction = prediction.reshape(prediction.shape[0], -1)
target = target.reshape(target.shape[0], -1)
numerator = torch.linalg.vector_norm(prediction - target, dim=1)
denominator = torch.linalg.vector_norm(target, dim=1).clamp_min(1.0e-8)
return numerator / denominator
|