Spaces:
Build error
Build error
File size: 1,209 Bytes
26246bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import logging
import typing as tp
import dora
import torch
logger = logging.getLogger(__name__)
class Profiler:
"""Context manager wrapper for xformers profiler.
"""
def __init__(self, module: torch.nn.Module, enabled: bool = False):
self.profiler: tp.Optional[tp.Any] = None
if enabled:
from xformers.profiler import profile
output_dir = dora.get_xp().folder / 'profiler_data'
logger.info("Profiling activated, results with be saved to %s", output_dir)
self.profiler = profile(output_dir=output_dir, module=module)
def step(self):
if self.profiler is not None:
self.profiler.step() # type: ignore
def __enter__(self):
if self.profiler is not None:
return self.profiler.__enter__() # type: ignore
def __exit__(self, exc_type, exc_value, exc_tb):
if self.profiler is not None:
return self.profiler.__exit__(exc_type, exc_value, exc_tb) # type: ignore
|