/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/fx/profiler/experimental/profiler_function/activation_function.py
35 строк
1 KB
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
from typing import Tuple import torch from ..registry import meta_profiler_function # TODO: different activation has different FLOPs count, currently unused. _multiplier = { torch.nn.functional.relu: 1, torch.nn.functional.prelu: 4, torch.nn.functional.sigmoid: 4, torch.nn.functional.tanh: 5, torch.nn.functional.leaky_relu: 3, torch.nn.functional.elu: 4, torch.nn.functional.relu6: 2, torch.nn.functional.gelu: 9, torch.nn.functional.hardswish: 5, torch.nn.functional.hardsigmoid: 4, } @meta_profiler_function.register(torch.nn.functional.leaky_relu) @meta_profiler_function.register(torch.nn.functional.elu) @meta_profiler_function.register(torch.nn.functional.gelu) @meta_profiler_function.register(torch.nn.functional.relu6) @meta_profiler_function.register(torch.nn.functional.prelu) @meta_profiler_function.register(torch.nn.functional.relu) @meta_profiler_function.register(torch.nn.functional.sigmoid) @meta_profiler_function.register(torch.nn.functional.tanh) @meta_profiler_function.register(torch.nn.functional.hardswish) @meta_profiler_function.register(torch.nn.functional.hardsigmoid) def torch_nn_func_non_linear_act(input: torch.Tensor, inplace: bool = False) -> Tuple[int, int]: flops = input.numel() macs = 0 return flops, macs