From ea14152ae65ca379584f255965f33bab51f05a5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Tue, 18 Apr 2023 20:29:09 +0200 Subject: [PATCH] Add PyTorch profiler to support tracing of forward/backward calls (#1761) * Add PyTorch profiler support * Add support for wrapped code, fix linting errors * Update docs --- .../debugging_tools/profiling.md | 27 +++- nerfstudio/configs/base_config.py | 8 +- nerfstudio/engine/trainer.py | 2 +- nerfstudio/utils/decorators.py | 2 +- nerfstudio/utils/profiler.py | 121 ++++++++++++++++-- 5 files changed, 142 insertions(+), 18 deletions(-) diff --git a/docs/developer_guides/debugging_tools/profiling.md b/docs/developer_guides/debugging_tools/profiling.md index ed8796cec..aabada5f4 100644 --- a/docs/developer_guides/debugging_tools/profiling.md +++ b/docs/developer_guides/debugging_tools/profiling.md @@ -23,7 +23,24 @@ class RayGenerator(nn.Module): ... ``` -At termination of training or end of the training run, the profiler will print out the average execution time for all of the functions that have the profiler tag. +Alternatively, you can also time parts of the code: +```python +... + +def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle: + # implementation here + with profiler.time_function("code1"): + # some code here + ... + + with profiler.time_function("code2"): + # some code here + ... + ... +``` + + +At termination of training or end of the training run, the profiler will print out the average execution time for all of the functions or code blocks that have the profiler tag. :::{admonition} Tip :class: info @@ -32,6 +49,14 @@ Use this profiler if there are *specific/individual functions* you want to measu ::: +#### Profiling with PyTorch profiler + +If you want to profile the training or evaluation code and track the memory and CUDA kernel launches, consider using [PyTorch profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html). +It will run the profiler for some selected step numbers, once with `CUDA_LAUNCH_BLOCKING=1`, once with `CUDA_LAUNCH_BLOCKING=0`. +The PyTorch profiler can be enabled with `--logging.profiler=pytorch` flag. +The outputs of the profiler are trace files stored in `{PATH_TO_MODEL_OUTPUT}/profiler_traces`, and can be loaded in Google Chrome by typing `chrome://tracing`. + + #### Profiling with PySpy If you want to profile the entire codebase, consider using [PySpy](https://github.com/benfred/py-spy). diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index fd5608b4e..b2665a649 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -119,9 +119,11 @@ class LoggingConfig(PrintableConfig): e.g. if 20, averages will be computed over past 20 occurrences.""" local_writer: LocalWriterConfig = LocalWriterConfig(enable=True) """if provided, will print stats locally. if None, will disable printing""" - enable_profiler: bool = True - """whether to enable profiling code; prints speed of functions at the end of a program. - profiler logs run times of functions and prints at end of training""" + profiler: Literal["none", "basic", "pytorch"] = "basic" + """how to profile the code; + "basic" - prints speed of all decorated functions at the end of a program. + "pytorch" - same as basic, but it also traces few training steps. + """ # Viewer related configs diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 7f07a3b5d..1dddb21d5 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -184,7 +184,7 @@ class Trainer: self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages ) writer.put_config(name="config", config_dict=dataclasses.asdict(self.config), step=0) - profiler.setup_profiler(self.config.logging) + profiler.setup_profiler(self.config.logging, writer_log_path) def setup_optimizers(self) -> Optimizers: """Helper to set up the optimizers diff --git a/nerfstudio/utils/decorators.py b/nerfstudio/utils/decorators.py index 6f92ac679..356369f82 100644 --- a/nerfstudio/utils/decorators.py +++ b/nerfstudio/utils/decorators.py @@ -42,7 +42,7 @@ def check_profiler_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.enable_profiler: + if self.config.profiler != "none": ret = func(self, *args, **kwargs) return ret diff --git a/nerfstudio/utils/profiler.py b/nerfstudio/utils/profiler.py index c1fc96dcb..246569444 100644 --- a/nerfstudio/utils/profiler.py +++ b/nerfstudio/utils/profiler.py @@ -17,10 +17,17 @@ Profiler base class and functionality """ from __future__ import annotations +import functools +import os import time -from typing import Callable +from collections import deque +from contextlib import contextmanager +from pathlib import Path +from typing import Callable, Dict, List, Optional, Tuple, Union from rich.console import Console +from torch.autograd.profiler import ContextDecorator +from torch.profiler import ProfilerActivity, profile, record_function from nerfstudio.configs import base_config as cfg from nerfstudio.utils import comms @@ -33,32 +40,122 @@ from nerfstudio.utils.decorators import ( CONSOLE = Console(width=120) PROFILER = [] +PYTORCH_PROFILER = None -def time_function(func: Callable) -> Callable: - """Decorator: time a function call""" +class time_function(ContextDecorator): # pylint: disable=invalid-name + """Decorator/Context manager: time a function call or a block of code""" - def wrapper(*args, **kwargs): - start = time.time() - ret = func(*args, **kwargs) + def __init__(self, name: str): + self.name = name + self.start = None + self._profiler_contexts = deque() + self._function_call_args: Optional[Tuple[Tuple, Dict]] = None + + def __new__(cls, func: Union[str, Callable]): + instance = super().__new__(cls) + if isinstance(func, str): + instance.__init__(func) + return instance + if callable(func): + instance.__init__(func.__qualname__) + return instance(func) + raise ValueError(f"Argument func of type {type(func)} is not a string or a callable.") + + def __enter__(self): + self.start = time.time() + if PYTORCH_PROFILER is not None: + args, kwargs = tuple(), {} + if self._function_call_args is not None: + args, kwargs = self._function_call_args + ctx = PYTORCH_PROFILER.record_function(self.name, *args, **kwargs) + ctx.__enter__() # pylint: disable=no-member + self._profiler_contexts.append(ctx) + if self._function_call_args is None: + ctx = record_function(self.name) + ctx.__enter__() + self._profiler_contexts.append(ctx) + + def __exit__(self, *args, **kwargs): + while self._profiler_contexts: + context = self._profiler_contexts.pop() + context.__exit__(*args, **kwargs) if PROFILER: - class_str = func.__qualname__ - PROFILER[0].update_time(class_str, start, time.time()) - return ret + PROFILER[0].update_time(self.name, self.start, time.time()) - return wrapper + def __call__(self, func: Callable): + @functools.wraps(func) + def inner(*args, **kwargs): + self._function_call_args = (args, kwargs) + with self: + out = func(*args, **kwargs) + self._function_call_args = None + return out + + return inner def flush_profiler(config: cfg.LoggingConfig): """Method that checks if profiler is enabled before flushing""" - if config.enable_profiler and PROFILER: + if config.profiler != "none" and PROFILER: PROFILER[0].print_profile() -def setup_profiler(config: cfg.LoggingConfig): +def setup_profiler(config: cfg.LoggingConfig, log_dir: Path): """Initialization of profilers""" + global PYTORCH_PROFILER # pylint: disable=global-statement if comms.is_main_process(): PROFILER.append(Profiler(config)) + if config.profiler == "pytorch": + PYTORCH_PROFILER = PytorchProfiler(log_dir) + + +class PytorchProfiler: # pylint: disable=too-few-public-methods + """ + Wrapper for Pytorch Profiler + """ + + def __init__(self, output_path: Path, trace_steps: Optional[List[int]] = None): + self.output_path = output_path / "profiler_traces" + if trace_steps is None: + # Some arbitrary steps which likely do not overlap with steps usually chosen to run callbacks + trace_steps = [12, 17] + self.trace_steps = trace_steps + + @contextmanager + def record_function(self, function: str, *args, **_kwargs): + """ + Context manager that records a function call and saves the trace to a json file. + Traced functions are: train_iteration, eval_iteration + """ + if function.endswith("train_iteration") or function.endswith("eval_iteration"): + step = args[1] + assert isinstance(step, int) + assert len(args) == 2 + stage = function.split(".")[-1].split("_")[0] + if step in self.trace_steps: + launch_kernel_blocking = self.trace_steps.index(step) % 2 == 0 + backup_lb_var = "" + if launch_kernel_blocking: + backup_lb_var = os.environ.get("CUDA_LAUNCH_BLOCKING", "") + os.environ["CUDA_LAUNCH_BLOCKING"] = "1" + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + with_stack=True, + profile_memory=True, + ) as prof: + yield None + if launch_kernel_blocking: + os.environ["CUDA_LAUNCH_BLOCKING"] = backup_lb_var + self.output_path.mkdir(parents=True, exist_ok=True) + prof.export_chrome_trace( + str(self.output_path / f"trace_{stage}_{step}{'_blocking' if launch_kernel_blocking else ''}.json") + ) + return + # Functions are recorded automatically + yield None + return @decorate_all([check_profiler_enabled, check_main_thread])