Inheritance diagram for fastreid.engine.hooks.AutogradProfiler:
Public Member Functions | |
| __init__ (self, enable_predicate, output_dir, *use_cuda=True) | |
| before_step (self) | |
| after_step (self) | |
 Public Member Functions inherited from fastreid.engine.train_loop.HookBase | |
| before_train (self) | |
| after_train (self) | |
Public Attributes | |
| trainer | |
Protected Attributes | |
| _enable_predicate | |
| _use_cuda | |
| _output_dir | |
| _profiler | |
Detailed Description
A hook which runs `torch.autograd.profiler.profile`.
Examples:
.. code-block:: python
hooks.AutogradProfiler(
lambda trainer: trainer.iter > 10 and trainer.iter < 20, self.cfg.OUTPUT_DIR
)
The above example will run the profiler for iteration 10~20 and dump
results to ``OUTPUT_DIR``. We did not profile the first few iterations
because they are typically slower than the rest.
The result files can be loaded in the ``chrome://tracing`` page in chrome browser.
Note:
When used together with NCCL on older version of GPUs,
autograd profiler may cause deadlock because it unnecessarily allocates
memory on every device it sees. The memory management calls, if
interleaved with NCCL calls, lead to deadlock on GPUs that do not
support `cudaLaunchCooperativeKernelMultiDevice`.
Constructor & Destructor Documentation
◆ __init__()
| fastreid.engine.hooks.AutogradProfiler.__init__ | ( | self, | |
| enable_predicate, | |||
| output_dir, | |||
| * | use_cuda = True ) |
Args:
enable_predicate (callable[trainer -> bool]): a function which takes a trainer,
and returns whether to enable the profiler.
It will be called once every step, and can be used to select which steps to profile.
output_dir (str): the output directory to dump tracing files.
use_cuda (bool): same as in `torch.autograd.profiler.profile`.
Definition at line 252 of file hooks.py.
252 def __init__(self, enable_predicate, output_dir, *, use_cuda=True):
253 """
254 Args:
255 enable_predicate (callable[trainer -> bool]): a function which takes a trainer,
256 and returns whether to enable the profiler.
257 It will be called once every step, and can be used to select which steps to profile.
258 output_dir (str): the output directory to dump tracing files.
259 use_cuda (bool): same as in `torch.autograd.profiler.profile`.
260 """
261 self._enable_predicate = enable_predicate
262 self._use_cuda = use_cuda
263 self._output_dir = output_dir
264
Member Function Documentation
◆ after_step()
| fastreid.engine.hooks.AutogradProfiler.after_step | ( | self | ) |
Called after each iteration.
Reimplemented from fastreid.engine.train_loop.HookBase.
Definition at line 272 of file hooks.py.
272 def after_step(self):
273 if self._profiler is None:
274 return
275 self._profiler.__exit__(None, None, None)
276 out_file = os.path.join(
277 self._output_dir, "profiler-trace-iter{}.json".format(self.trainer.iter)
278 )
279 if "://" not in out_file:
280 self._profiler.export_chrome_trace(out_file)
281 else:
282 # Support non-posix filesystems
283 with tempfile.TemporaryDirectory(prefix="fastreid_profiler") as d:
284 tmp_file = os.path.join(d, "tmp.json")
285 self._profiler.export_chrome_trace(tmp_file)
286 with open(tmp_file) as f:
287 content = f.read()
288 with PathManager.open(out_file, "w") as f:
289 f.write(content)
290
291
◆ before_step()
| fastreid.engine.hooks.AutogradProfiler.before_step | ( | self | ) |
Called before each iteration.
Reimplemented from fastreid.engine.train_loop.HookBase.
Definition at line 265 of file hooks.py.
265 def before_step(self):
266 if self._enable_predicate(self.trainer):
267 self._profiler = torch.autograd.profiler.profile(use_cuda=self._use_cuda)
268 self._profiler.__enter__()
269 else:
270 self._profiler = None
271
Member Data Documentation
◆ _enable_predicate
|
protected |
◆ _output_dir
|
protected |
◆ _profiler
◆ _use_cuda
◆ trainer
The documentation for this class was generated from the following file:
- smreid/fastreid/engine/hooks.py
