/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
examples/language/model_utils.py
32 строки
713 B
flybird11111
[example]add gpt2 benchmark example script. (#5295)
04 мар 2024, 11:18
Не верифицирован
04 мар 2024, 11:18
29695cf
Код
Авторство
О чём код?
from contextlib import contextmanager import torch import torch.nn as nn @contextmanager def low_precision_init(target_dtype: torch.dtype = torch.float16): dtype = torch.get_default_dtype() try: torch.set_default_dtype(target_dtype) yield finally: torch.set_default_dtype(dtype) def get_model_numel(model: nn.Module) -> int: return sum(p.numel() for p in model.parameters()) def format_numel_str(numel: int) -> str: B = 1024**3 M = 1024**2 K = 1024 if numel >= B: return f"{numel / B:.2f} B" elif numel >= M: return f"{numel / M:.2f} M" elif numel >= K: return f"{numel / K:.2f} K" else: return f"{numel}"