# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils.ipynb.

# %% auto 0
__all__ = ['free_gpu_memory', 'less_random', 'scale_time']

# %% ../../nbs/utils.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../../nbs/utils.ipynb 3
import torch, random, gc
import numpy as np
from fastai.learner import Learner
from fastai.data.core import DataLoaders
from fastcore.foundation import contextmanager
from fastai.callback.core import set_random_states, get_random_states

# %% ../../nbs/utils.ipynb 4
def free_gpu_memory(learn:Learner, dls:DataLoaders=None):
    "Frees GPU memory using `gc.collect` and `torch.cuda.empty_cache`"
    learn.dls, learn, dls = None, None, None
    gc.collect()
    torch.cuda.empty_cache()

# %% ../../nbs/utils.ipynb 5
@contextmanager
def less_random(seed=42, deterministic=True):
    "Stores and retrieves state of random number generators. Sets random seed for `random`, `torch`, and `numpy`. Does not set `torch.backends.cudnn.benchmark = False`"
    states = get_random_states()

    try: torch.manual_seed(seed)
    except NameError: pass
    try: torch.cuda.manual_seed_all(seed)
    except NameError: pass
    try: np.random.seed(seed%(2**32-1))
    except NameError: pass
    random.seed(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = True
    try:
        yield #we are managing global variables
    finally:
        set_random_states(**states)

# %% ../../nbs/utils.ipynb 7
# modified from https://github.com/thomasbrandon/mish-cuda/blob/master/test/perftest.py
def scale_time(val:float, spec:str="#0.4G"):
    "Scale fractional second `time` values and return formatted to `spec`"
    if val == 0: return '-'
    PREFIXES = np.array([c for c in u"yzafpnµm kMGTPEZY"])
    exp = np.int8(np.log10(np.abs(val)) // 3 * 3 * np.sign(val))
    val /= 10.**exp
    prefix = PREFIXES[exp//3 + len(PREFIXES)//2]
    return f"{val:{spec}}{prefix}s"
