27 lines
517 B
Python
27 lines
517 B
Python
import dataclasses
|
|
from typing import Callable, Optional
|
|
|
|
|
|
all_experiments: dict[str, Callable] = {}
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Experiment:
|
|
name: str
|
|
metric: str
|
|
target: float
|
|
actual: float
|
|
dtype: str
|
|
device: str
|
|
arch: str # GPU name for CUDA or CPU arch for CPU
|
|
is_model: bool = False
|
|
|
|
|
|
def register_experiment(name: Optional[str] = None):
|
|
def decorator(func):
|
|
key = name or func.__name__
|
|
all_experiments[key] = func
|
|
return func
|
|
|
|
return decorator
|