/
vshmidt
/
pytorch
Обзор
Документация
Войти
/
vshmidt
/
pytorch
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
functorch/examples/compilation/simple_function.py
35 строк
761 B
Xuehai Pan
[BE][Easy][4/19] enforce style for empty lines in import segments in `functorch/` (#129755)
18 июл 2024, 08:08
18 июл 2024, 08:08
740fb22
Код
Авторство
О чём код?
# Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import time import torch from functorch import grad, make_fx from functorch.compile import nnc_jit def f(x): return torch.sin(x).sum() inp = torch.randn(100) grad_pt = grad(f) grad_fx = make_fx(grad_pt)(inp) grad_nnc = nnc_jit(grad_pt) def bench(name, f, iters=10000, warmup=3): for _ in range(warmup): f() begin = time.time() for _ in range(iters): f() print(f"{name}: ", time.time() - begin) bench("Pytorch: ", lambda: grad_pt(inp)) bench("FX: ", lambda: grad_fx(inp)) bench("NNC: ", lambda: grad_nnc(inp))