torch_utils/utils/memory.py
2021-02-25 02:18:02 +09:00

27 lines
912 B
Python

def human_size(byte_count: int) -> str:
"""Output byte amount in human readable format"""
amount = float(byte_count)
for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi', 'Yi']:
if amount < 1024.0:
break
amount /= 1024.0
return f'{amount:.2f}{unit}B'
def human_to_bytes(text: str) -> float:
split_index = 0
while '0' <= text[split_index] <= '9':
split_index += 1
if split_index == len(text):
return float(text)
amount = float(text[:split_index])
unit = text[split_index:].strip()
if not unit:
return amount
if unit not in ['KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB']:
raise RuntimeError(f'Unrecognized unit : {unit}')
for final_unit in ['KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB']:
amount *= 1024.0
if unit == final_unit:
return amount