/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/fx/profiler/experimental/profiler_function/linear.py
15 строк
438 B
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 @meta_profiler_function.register(torch.nn.functional.linear) def torch_nn_linear(input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor = None) -> Tuple[int, int]: out_features = weight.shape[0] macs = torch.numel(input) * out_features flops = 2 * macs if bias is not None: flops += bias.numel() return flops, macs