Python get gpu memory
14 Python code examples are found related to "
get gpu memory".
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
Example 1
Source File: memory.py From pytorch-lightning with Apache License 2.0 | 7 votes |
def get_gpu_memory_map() -> Dict[str, int]: """Get the current gpu usage. Return: A dictionary in which the keys are device ids as integers and values are memory usage as integers in MB. """ result = subprocess.run( ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader",], encoding="utf-8", # capture_output=True, # valid for python version >=3.7 stdout=PIPE, stderr=PIPE, # for backward compatibility with python version 3.6 check=True, ) # Convert lines into a dictionary gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)] gpu_memory_map = {f"gpu_{index}": memory for index, memory in enumerate(gpu_memory)} return gpu_memory_map
Example 2
Source File: utils.py From pytorch-apex-experiment with MIT License | 6 votes |
def get_gpu_memory_map(): """Get the current gpu usage. Returns ------- usage: dict Keys are device ids as integers. Values are memory usage as integers in MB. """ result = subprocess.check_output( [ 'nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader' ]) return int(result)
Example 3
Source File: utils.py From sockeye with Apache License 2.0 | 6 votes |
def get_gpu_memory_usage(ctx: Union[mx.context.Context, List[mx.context.Context]]) -> Dict[int, Tuple[int, int]]: """ Returns used and total memory for GPUs identified by the given context list. :param ctx: List of MXNet context devices. :return: Dictionary of device id mapping to a tuple of (memory used, memory total). """ if isinstance(ctx, mx.context.Context): ctx = [ctx] ctx = [c for c in ctx if c.device_type == 'gpu'] if not ctx: return {} memory_data = {} # type: Dict[int, Tuple[int, int]] for c in ctx: try: free, total = mx.context.gpu_memory_info(device_id=c.device_id) # in bytes used = total - free memory_data[c.device_id] = (used * 1e-06, total * 1e-06) except mx.MXNetError: logger.exception("Failed retrieving memory data for gpu%d", c.device_id) continue log_gpu_memory_usage(memory_data) return memory_data
Example 4
Source File: utils.py From gobbli with Apache License 2.0 | 5 votes |
def get_gpu_memory_map(): result = subprocess.check_output( [ 'nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader' ], encoding='utf-8') gpu_memory = [int(x) for x in result.strip().split('\n')] gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory)) return gpu_memory_map
Example 5
Source File: utils.py From DLInfBench with Apache License 2.0 | 5 votes |
def get_gpu_memory(pid=os.getpid()): info = os.popen('nvidia-smi').read() for line in info.split('\n'): t = line.strip().split() try: cur_pid = int(t[2]) if cur_pid != pid: continue gpu_mem = int(t[5][:-3]) return gpu_mem except: continue return 0
Example 6
Source File: main.py From joint-cnn-mrf with MIT License | 5 votes |
def get_gpu_memory(gpus): """ Small heuristic to calculate the amount of memory needed for each GPU in case of multi-gpu training. """ if len(gpus) >= 5: return 0.4 elif len(gpus) >= 3: return 0.5 elif len(gpus) == 2: return 0.6 else: return 0.6
Example 7
Source File: utils.py From sg2im with Apache License 2.0 | 5 votes |
def get_gpu_memory(): torch.cuda.synchronize() opts = [ 'nvidia-smi', '-q', '--gpu=' + str(0), '|', 'grep', '"Used GPU Memory"' ] cmd = str.join(' ', opts) ps = subprocess.Popen(cmd,shell=True,stdout=subprocess.PIPE,stderr=subprocess.STDOUT) output = ps.communicate()[0].decode('utf-8') output = output.split("\n")[1].split(":") consumed_mem = int(output[1].strip().split(" ")[0]) return consumed_mem
Example 8
Source File: utils.py From MT-DNN with MIT License | 5 votes |
def get_gpu_memory_map(): result = subprocess.check_output( ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"], encoding="utf-8", ) gpu_memory = [int(x) for x in result.strip().split("\n")] gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory)) return gpu_memory_map
Example 9
Source File: like_top.py From bifrost with BSD 3-Clause "New" or "Revised" License | 5 votes |
def get_gpu_memory_usage(): """ Grab nvidia-smi output and return a dictionary of the memory usage. """ data = {'devCount':0, 'memTotal':0, 'memUsed':0, 'memFree':0, 'pwrDraw':0.0, 'pwrLimit':0.0, 'load':0.0} q_flag = '--query-gpu=memory.used,memory.total,memory.free,power.draw,power.limit,utilization.gpu' fmt_flag = '--format=csv,noheader,nounits' try: p = subprocess.Popen(['nvidia-smi', q_flag, fmt_flag], stdout=subprocess.PIPE, stderr=subprocess.PIPE) output, err = p.communicate() except (OSError, ValueError) as e: pass else: # Parse the ouptut and turn everything into something useful, if possible lines = output.split('\n')[:-1] for line in lines: used, total, free, draw, limit, load = line.split(',') data['devCount'] += 1 data['memTotal'] += int(total, 10)*1024 data['memUsed'] += int(used, 10)*1024 data['memFree'] += int(free, 10)*1024 try: data['pwrDraw'] += float(draw) data['pwrLimit'] += float(limit) except ValueError: pass try: data['load'] += float(load) except ValueError: pass # Convert the load to an average data['load'] /= data['devCount'] return data
Example 10
Source File: utils.py From sgan with MIT License | 5 votes |
def get_gpu_memory(): torch.cuda.synchronize() opts = [ 'nvidia-smi', '-q', '--gpu=' + str(1), '|', 'grep', '"Used GPU Memory"' ] cmd = str.join(' ', opts) ps = subprocess.Popen( cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = ps.communicate()[0].decode('utf-8') output = output.split("\n")[0].split(":") consumed_mem = int(output[1].strip().split(" ")[0]) return consumed_mem
Example 11
Source File: find.py From second.pytorch with MIT License | 4 votes |
def get_gpu_memory_usage(): if sys.platform == 'win32': # TODO: add windows support return None cuda_home = find_cuda() if cuda_home is None: return None cuda_home = Path(cuda_home) source = """ #include <cuda_runtime.h> #include <iostream> int main(){ int nDevices; cudaGetDeviceCount(&nDevices); size_t free_m, total_m; // output json format. std::cout << "["; for (int i = 0; i < nDevices; i++) { cudaSetDevice(i); cudaMemGetInfo(&free_m, &total_m); std::cout << "[" << free_m << "," << total_m << "]"; if (i != nDevices - 1) std::cout << "," << std::endl; } std::cout << "]" << std::endl; return 0; } """ with tempfile.NamedTemporaryFile('w', suffix='.cc') as f: f_path = Path(f.name) f.write(source) f.flush() try: # TODO: add windows support cmd = ( f"g++ {f.name} -o {f_path.stem} -std=c++11" f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart") print(cmd) subprocess.check_output(cmd, shell=True, cwd=f_path.parent) cmd = f"./{f_path.stem}" usages = subprocess.check_output( cmd, shell=True, cwd=f_path.parent).decode() usages = json.loads(usages) return usages except: return None return None
Example 12
Source File: find.py From Det3D with Apache License 2.0 | 4 votes |
def get_gpu_memory_usage(): if sys.platform == "win32": # TODO: add windows support return None cuda_home = find_cuda() if cuda_home is None: return None cuda_home = Path(cuda_home) source = """ #include <cuda_runtime.h> #include <iostream> int main(){ int nDevices; cudaGetDeviceCount(&nDevices); size_t free_m, total_m; // output json format. std::cout << "["; for (int i = 0; i < nDevices; i++) { cudaSetDevice(i); cudaMemGetInfo(&free_m, &total_m); std::cout << "[" << free_m << "," << total_m << "]"; if (i != nDevices - 1) std::cout << "," << std::endl; } std::cout << "]" << std::endl; return 0; } """ with tempfile.NamedTemporaryFile("w", suffix=".cc") as f: f_path = Path(f.name) f.write(source) f.flush() try: # TODO: add windows support cmd = ( f"g++ {f.name} -o {f_path.stem} -std=c++11" f" -I{cuda_home / 'include'} -L{cuda_home / 'lib64'} -lcudart" ) print(cmd) subprocess.check_output(cmd, shell=True, cwd=f_path.parent) cmd = f"./{f_path.stem}" usages = subprocess.check_output( cmd, shell=True, cwd=f_path.parent ).decode() usages = json.loads(usages) return usages except Exception: return None return None