PyTorchのプロファイラ

PyTorch Lightning:

Class	device	unit	internal
SimpleProfiler		L's action
PyTorchProfiler			PyTorch’s Autograd Profiler
AdvancedProfiler
XLAProfiler	TPU
(PassThroughProfiler)	any		(default) Just pass-through, do nothing

SimpleProfilerとDataLoader

データロードの経過時間は _TrainingEpochLoop.train_dataloader_next アクションとして記録される。
worker有りのDataLoaderはprefetchをするため、理想的にはこのアクションがほぼ0になる。
ここが非ゼロということはデータ供給が追い付いていないことを意味する。

Trainerの奥深くにしまい込まれており、実装の理解は面倒。
call関係は以下になっており、Trainer.fitの前処理でprofilerがフックされている。

<Trainer>.__init__(...):
    self.fit_loop = _FitLoop(...)

<Trainer>.fit(...):
    call._call_and_handle_interrupt(self._fit_impl, ...)

<Trainer>._fit_impl(...):
    self._run(model, ...)

<Trainer>._run(model, ...):
    results = self._run_stage()

<Trainer>._run_stage():
    self.fit_loop.run()


<_FitLoop>.__init__(...):
    self.epoch_loop = _TrainingEpochLoop(trainer)

<_FitLoop>.run():
    while not self.done:
        self.advance()

<_FitLoop>.advance():
    self.epoch_loop.run(self._data_fetcher)


<_TrainingEpochLoop>.run(data_fetcher):
    self.on_run_start(data_fetcher)
    while not self.done:
        self.advance(data_fetcher)

<_TrainingEpochLoop>.on_run_start(data_fetcher):
    data_fetcher._start_profiler = self._on_before_fetch
    data_fetcher._stop_profiler  = self._on_after_fetch
<_TrainingEpochLoop>._on_before_fetch():
    self.trainer.profiler.start(f"[{self.__class__.__name__}].train_dataloader_next")
<_TrainingEpochLoop>._on_after_fetch():
    self.trainer.profiler.stop(f"[{self.__class__.__name__}].train_dataloader_next")

<_TrainingEpochLoop>.advance(data_fetcher):
    next(data_fetcher)


<_DataFetcher>.__next__():
    self._start_profiler()
    data = next(self.iterator)
    self._stop_profiler()
    return data

<_PrefetchDataFetcher>.__next__():
    batch = self.batches.pop(0)
    self._fetch_next_batch(self.iterator)
    return batch

<_PrefetchDataFetcher>._fetch_next_batch(iterator):
    self._start_profiler()
    batch = next(iterator)
    self._stop_profiler()
    self.batches.append(batch)

_data_fetcherの初期化に関して:

<_FitLoop>.__init__(...):
    self._combined_loader = None
    self._data_fetcher = None

<_FitLoop>.run():
    self.setup_data()
    self.reset()
    self.on_run_start()
    while not self.done:
        self.on_advance_start()
        self.advance()
        self.on_advance_end()
    self.on_run_end()

<_FitLoop>.on_run_start():
    self._data_fetcher = _select_data_fetcher(trainer)

<_FitLoop>.advance():
    combined_loader = self._combined_loader
    self._data_fetcher.setup(combined_loader)
    self.epoch_loop.run(self._data_fetcher)


def _select_data_fetcher(trainer) -> _DataFetcher:
    lightning_module = trainer.lightning_module
    if trainer.testing:
        step_fx_name = "test_step"
    elif trainer.training:
        step_fx_name = "training_step"
    elif trainer.validating or trainer.sanity_checking:
        step_fx_name = "validation_step"
    elif trainer.predicting:
        step_fx_name = "predict_step"
    else:
        raise RuntimeError(f"DataFetcher is unsupported for {trainer.state.stage}")
    step_fx = getattr(lightning_module, step_fx_name)
    if is_param_in_hook_signature(step_fx, "dataloader_iter", explicit=True):
        rank_zero_warn(
            f"Found `dataloader_iter` argument in the `{step_fx_name}`. Note that the support for "
            "this signature is experimental and the behavior is subject to change."
        )
        return _DataLoaderIterDataFetcher()
    return _PrefetchDataFetcher()