File size: 1,475 Bytes
27d37fd | 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 | from typing import Union
import json
import os
import random
from pathlib import Path
import numpy as np
import torch
def ensure_dir(path: Union[str, Path]) -> None:
"""
Input:
path: Directory path.
Output:
Creates the directory if it does not exist.
"""
Path(path).mkdir(parents=True, exist_ok=True)
def save_json(obj: dict, path: Union[str, Path]) -> None:
"""
Input:
obj: Python dictionary.
path: Output JSON path.
Output:
Saves dictionary as formatted JSON.
"""
with open(path, "w", encoding="utf-8") as f:
json.dump(obj, f, indent=2, ensure_ascii=False)
def load_json(path: Union[str, Path]) -> dict:
"""
Input:
path: JSON file path.
Output:
Loaded Python dictionary.
"""
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def set_seed(seed: int) -> None:
"""
Input:
seed: Integer seed.
Output:
Makes Python, NumPy, and PyTorch randomness more reproducible.
"""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def is_rank_zero() -> bool:
"""
Output:
True only for the main distributed process.
"""
return int(os.environ.get("RANK", "0")) == 0
def get_world_size() -> int:
"""
Output:
Number of distributed processes.
"""
return int(os.environ.get("WORLD_SIZE", "1"))
|