Module audiocraft.utils.profiler

Classes

class Profiler (module: torch.nn.modules.module.Module, enabled: bool = False)

Expand source code

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

Context manager wrapper for xformers profiler.

Methods

def step(self)

Expand source code

def step(self):
    if self.profiler is not None:
        self.profiler.step()  # type: ignore