/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/fx/profiler/experimental/profiler_module/normalization.py
43 строки
2 KB
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
# Copyright (c) Microsoft Corporation. # Licensed under the MIT License. from typing import Tuple, Union import torch from ..registry import meta_profiler_module @meta_profiler_module.register(torch.nn.InstanceNorm1d) @meta_profiler_module.register(torch.nn.InstanceNorm2d) @meta_profiler_module.register(torch.nn.InstanceNorm3d) @meta_profiler_module.register(torch.nn.LayerNorm) @meta_profiler_module.register(torch.nn.GroupNorm) @meta_profiler_module.register(torch.nn.BatchNorm1d) @meta_profiler_module.register(torch.nn.BatchNorm2d) @meta_profiler_module.register(torch.nn.BatchNorm3d) def torch_nn_normalize( self: Union[ torch.nn.LayerNorm, torch.nn.GroupNorm, torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, torch.nn.BatchNorm3d ], input: torch.Tensor, ) -> Tuple[int, int]: # adopted from https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/profiling/flops_profiler/profiler.py#L615 has_affine = self.weight is not None if self.training: flops = input.numel() * (2 if has_affine else 1) else: flops = input.numel() * (5 if has_affine else 4) macs = 0 return flops, macs try: import apex meta_profiler_module.register(apex.normalization.FusedLayerNorm)(torch_nn_normalize) meta_profiler_module.register(apex.normalization.FusedRMSNorm)(torch_nn_normalize) meta_profiler_module.register(apex.normalization.MixedFusedLayerNorm)(torch_nn_normalize) meta_profiler_module.register(apex.normalization.MixedFusedRMSNorm)(torch_nn_normalize) except (ImportError, AttributeError): pass