Source code for draugr.torch_utilities.system.info

import multiprocessing
import platform
import typing

import torch

__all__ = ["system_info", "cuda_info"]


[docs]def system_info() -> str: """ :return: :rtype:""" return "\n".join( [ f"Python version: {platform.python_version()}", f"Python implementation: {platform.python_implementation()}", f"Python compiler: {platform.python_compiler()}", f"PyTorch version: {torch.__version__}", f"System: {platform.system() or 'Unable to determine'}", f"System version: {platform.release() or 'Unable to determine'}", f"Processor: {platform.processor() or 'Unable to determine'}", f"Number of CPUs: {multiprocessing.cpu_count()}", ] )
[docs]def cuda_info() -> str: """ :return: :rtype:""" def _cuda_devices_formatting( info_function: typing.Callable, formatting_function: typing.Callable = None, mapping_function: typing.Callable = None, ): def _setup_default(function): return (lambda arg: arg) if function is None else function formatting_function = _setup_default(formatting_function) mapping_function = _setup_default(mapping_function) return " | ".join( mapping_function( [ formatting_function(info_function(i)) for i in range(torch.cuda.device_count()) ] ) ) def _device_properties(attribute): return _cuda_devices_formatting( lambda i: getattr(torch.cuda.get_device_properties(i), attribute), mapping_function=lambda in_bytes: map(str, in_bytes), ) cuda_cap = _cuda_devices_formatting( torch.cuda.get_device_capability, formatting_function=lambda capabilities: ".".join(map(str, capabilities)), ) return "\n".join( [ f"Available CUDA devices count: {torch.cuda.device_count()}", f"CUDA devices names: {_cuda_devices_formatting(torch.cuda.get_device_name)}", f"Major.Minor CUDA capabilities of devices: {cuda_cap}", f"Device total memory (bytes): {_device_properties('total_memory')}", f"Device multiprocessor count: {_device_properties('multi_processor_count')}", ] )
if __name__ == "__main__": print(system_info()) print(cuda_info())