mirror of https://github.com/hpcaitech/ColossalAI
44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# 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
|