ThroughputMonitor

class lightning.pytorch.callbacks.ThroughputMonitor(batch_size_fn, length_fn=None, **kwargs)

Bases: Callback

Computes and logs throughput with the Throughput

Example:

class MyModel(LightningModule):
    def setup(self, stage):
        with torch.device("meta"):
            model = MyModel()

            def sample_forward():
                batch = torch.randn(..., device="meta")
                return model(batch)

            self.flops_per_batch = measure_flops(model, sample_forward, loss_fn=torch.Tensor.sum)


logger = ...
throughput = ThroughputMonitor(batch_size_fn=lambda batch: batch.size(0))
trainer = Trainer(max_steps=1000, log_every_n_steps=10, callbacks=throughput, logger=logger)
model = MyModel()
trainer.fit(model)

Notes

• It assumes that the batch size is the same during all iterations.

• It will try to access a flops_per_batch attribute on your LightningModule on every iteration. We suggest using the measure_flops() function for this. You might want to compute it differently each time based on your setup.

Parameters:

• batch_size_fn (Callable[[Any], int]) – A function to compute the number of samples given a batch.

• length_fn (Optional[Callable[[Any], int]]) – A function to compute the number of items in a sample given a batch.

• **kwargs (Any) – See available parameters in Throughput

on_predict_batch_end(trainer, pl_module, outputs, batch, *_, **__)

Called when the predict batch ends.

Return type:

None

on_predict_start(trainer, *_)

Called when the predict begins.

Return type:

None

on_test_batch_end(trainer, pl_module, outputs, batch, *_, **__)

Called when the test batch ends.

Return type:

None

on_test_start(trainer, *_)

Called when the test begins.

Return type:

None

on_train_batch_end(trainer, pl_module, outputs, batch, *_)

Called when the train batch ends. :rtype: None

Note

The value outputs["loss"] here will be the normalized value w.r.t accumulate_grad_batches of the loss returned from training_step.

on_train_start(trainer, *_)

Called when the train begins.

Return type:

None

on_validation_batch_end(trainer, pl_module, outputs, batch, *_, **__)

Called when the validation batch ends.

Return type:

None

on_validation_end(trainer, *_)

Called when the validation loop ends.

Return type:

None

on_validation_start(trainer, *_)

Called when the validation loop begins.

Return type:

None

setup(trainer, pl_module, stage)

Called when fit, validate, test, predict, or tune begins.

Return type:

None