saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refine machine_info and output onnxruntime_tools version (#4679)

Browse source 
* output onnxruntime_tools version
* change get_machine_info return data type to string
This commit is contained in:
Tianlei Wu 2020-08-02 18:20:59 -07:00 committed by GitHub
parent 6958f49dae
commit e70e9e2f67
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 34 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
onnxruntime/python/tools/transformers/machine_info.py
View file

@ -39,14 +39,16 @@ class MachineInfo():
def get_machine_info(self):
"""Get machine info in metric format"""
gpu_info = self.get_gpu_info_by_nvml()
cpu_info = cpuinfo.get_cpu_info()
machine_info = {
"gpu": gpu_info,
"cpu": self.get_cpu_info(),
"memory": self.get_memory_info(),
"python": cpuinfo.get_cpu_info()["python_version"], #sys.version,
"python": self._try_get(cpu_info, ["python_version"]),
"os": platform.platform(),
"onnxruntime": self.get_onnxruntime_info(),
"onnxruntime_tools": self.get_onnxruntime_tools_info(),
"pytorch": self.get_pytorch_info(),
"tensorflow": self.get_tensorflow_info()
}
@ -57,16 +59,23 @@ class MachineInfo():
mem = psutil.virtual_memory()
return {"total": mem.total, "available": mem.available}
def _try_get(self, cpu_info: Dict, names: List) -> str:
for name in names:
if name in cpu_info:
return cpu_info[name]
return ""
def get_cpu_info(self) -> Dict:
"""Get CPU info"""
cpu_info = cpuinfo.get_cpu_info()
return {
"brand": cpu_info["brand"] if "brand" in cpu_info else cpu_info["brand_raw"],
"brand": self._try_get(cpu_info, ["brand", "brand_raw"]),
"cores": psutil.cpu_count(logical=False),
"logical_cores": psutil.cpu_count(logical=True),
"hz": cpu_info["hz_actual"],
"l2_cache": cpu_info["l2_cache_size"],
"l3_cache": cpu_info["l3_cache_size"],
"hz": self._try_get(cpu_info, ["hz_actual"]),
"l2_cache": self._try_get(cpu_info, ["l2_cache_size"]),
"flags": self._try_get(cpu_info, ["flags"]),
"processor": platform.uname().processor
}
@ -114,10 +123,23 @@ class MachineInfo():
self.logger.exception(exception, False)
return None
def get_onnxruntime_tools_info(self) -> Dict:
try:
import onnxruntime_tools
return {"version": onnxruntime_tools.__version__}
except ImportError as error:
if not self.silent:
self.logger.exception(error)
return None
except Exception as exception:
if not self.silent:
self.logger.exception(exception, False)
return None
def get_pytorch_info(self) -> Dict:
try:
import torch
return {"version": torch.__version__, "support_gpu": torch.cuda.is_available()}
return {"version": torch.__version__, "support_gpu": torch.cuda.is_available(), "cuda": torch.version.cuda}
except ImportError as error:
if not self.silent:
self.logger.exception(error)
@ -155,7 +177,11 @@ def parse_arguments():
return args
def get_machine_info(silent=True) -> str:
machine = MachineInfo(silent)
return json.dumps(machine.machine_info, indent=2)
if __name__ == '__main__':
args = parse_arguments()
machine = MachineInfo(args.silent)
print(json.dumps(machine.machine_info, indent=2))
print(get_machine_info(args.silent))