"Helper utility functions for memory management"
from ipyexperiments.utils.pynvml_gate import load_pynvml_env
import gc
from collections import namedtuple
try:
import torch # currently relying on pytorch
except Exception as e:
raise Exception(f"{e}\nYou need to install the torch module; pip install torch")
use_gpu = torch.cuda.is_available()
if not use_gpu:
raise RuntimeError("these functions require CUDA environment; torch.cuda.is_available() returns false")
pynvml = load_pynvml_env()
############# gpu memory helper functions ############
GPUMemory = namedtuple('GPUMemory', ['total', 'free', 'used'])
def preload_pytorch():
""" Do a small operation on CUDA to get the pytorch/cuda structures in place.
A must to be run first if you're going to compare any CUDA-related numbers.
"""
torch.ones((1, 1)).cuda()
preload_pytorch() # needed to run first to get the measurements right
def b2mb(num):
def backend_init(self):
super().backend_init()
from ipyexperiments.utils.pynvml_gate import load_pynvml_env
self.pynvml = load_pynvml_env()