/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/fx/profiler/experimental/profiler_function/pooling.py
24 строки
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 @meta_profiler_function.register(torch.nn.functional.avg_pool1d) @meta_profiler_function.register(torch.nn.functional.avg_pool2d) @meta_profiler_function.register(torch.nn.functional.avg_pool3d) @meta_profiler_function.register(torch.nn.functional.max_pool1d) @meta_profiler_function.register(torch.nn.functional.max_pool2d) @meta_profiler_function.register(torch.nn.functional.max_pool3d) @meta_profiler_function.register(torch.nn.functional.adaptive_avg_pool1d) @meta_profiler_function.register(torch.nn.functional.adaptive_avg_pool2d) @meta_profiler_function.register(torch.nn.functional.adaptive_avg_pool3d) @meta_profiler_function.register(torch.nn.functional.adaptive_max_pool1d) @meta_profiler_function.register(torch.nn.functional.adaptive_max_pool2d) @meta_profiler_function.register(torch.nn.functional.adaptive_max_pool3d) def torch_nn_func_pooling(input: torch.Tensor, *args, **kwargs) -> Tuple[int, int]: # all pooling could be considered as going over each input element only once (https://stackoverflow.com/a/67301217) flops = input.numel() macs = 0 return flops, macs