otx.algo.callbacks.gpu_mem_monitor#

Monitor GPU memory hook.

Classes

GPUMemMonitor()

Monitor GPU memory hook.
class otx.algo.callbacks.gpu_mem_monitor.GPUMemMonitor[source]#

Bases: Callback

Monitor GPU memory hook.

on_train_batch_start(trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int) → None[source]#

Log GPU memory usage at the start of every train batch.

Parameters:

  • trainer (Trainer) – pl trainer.

  • pl_module (LightningModule) – pl module.

  • batch (Any) – current batch.

  • batch_idx (int) – current batch index.

on_validation_batch_start(trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int, dataloader_idx: int = 0) → None[source]#

Log GPU memory usage at the start of every validation batch.

Parameters:

  • trainer (Trainer) – pl trainer.

  • pl_module (LightningModule) – pl module.

  • batch (Any) – current batch.

  • batch_idx (int) – current batch index.

  • dataloader_idx (int, optional) – dataloader index. Defaults to 0.