File size: 2,560 Bytes
0d2f59e | 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 | """Configuration, reproducible IO and hardware checks shared by entrypoints."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
def arguments(description):
parser = argparse.ArgumentParser(description=description)
parser.add_argument('--config', type=Path, default=ROOT / 'config/config.yaml')
return parser
def load_config(path):
path = Path(path).resolve()
with path.open() as handle:
config = yaml.safe_load(handle)
return config, path.parents[1]
def output_path(root, value):
value = Path(value)
return value if value.is_absolute() else root / value
def save_json(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(path.suffix + '.tmp')
tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + '\n')
tmp.replace(path)
def sha256(path):
digest = hashlib.sha256()
with Path(path).open('rb') as handle:
for part in iter(lambda: handle.read(1024 * 1024), b''):
digest.update(part)
return digest.hexdigest()
def device_for(config, require_dcu=True):
device = torch.device(config['training']['device'])
if require_dcu:
if not os.environ.get('SLURM_JOB_ID'):
raise RuntimeError('This experiment must run in a SLURM DCU allocation.')
if device.type != 'cuda' or not torch.cuda.is_available() or not torch.version.hip:
raise RuntimeError('A HIP/DCU PyTorch device is required; CPU fallback is disabled.')
return device
def runtime_info(device):
return {'hostname': os.uname().nodename, 'slurm_job_id': os.environ.get('SLURM_JOB_ID'),
'slurm_partition': os.environ.get('SLURM_JOB_PARTITION'),
'torch': torch.__version__, 'hip': torch.version.hip,
'device': str(device), 'device_name': torch.cuda.get_device_name(device),
'device_memory_bytes': torch.cuda.get_device_properties(device).total_memory}
def errors(prediction, target):
difference = np.asarray(prediction, dtype=np.float64) - np.asarray(target, dtype=np.float64)
return {'MAE': np.mean(np.abs(difference), axis=0).tolist(),
'RMSE': np.sqrt(np.mean(difference ** 2, axis=0)).tolist()}
def checkpoint_save(path, payload):
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix('.tmp')
torch.save(payload, temporary)
temporary.replace(path)
|