batch_size_finder() — pytorch Function Reference

Architecture documentation for the batch_size_finder() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 2 called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ batch_size_finder() — pytorch Function Reference

Dependency Diagram

graph TD
  ba74afb9_ef25_84a8_66e2_1d25dcee421f["batch_size_finder()"]
  c9be2096_e6d7_2374_ad2e_a6e33f435ada["run()"]
  c9be2096_e6d7_2374_ad2e_a6e33f435ada -->|calls| ba74afb9_ef25_84a8_66e2_1d25dcee421f
  88e41ee6_6f08_0c6a_8717_6ba30d9b1ad4["decay_batch_exp()"]
  ba74afb9_ef25_84a8_66e2_1d25dcee421f -->|calls| 88e41ee6_6f08_0c6a_8717_6ba30d9b1ad4
  06ff896d_4db0_aa47_b3cd_be1da620f0ea["empty_gpu_cache()"]
  ba74afb9_ef25_84a8_66e2_1d25dcee421f -->|calls| 06ff896d_4db0_aa47_b3cd_be1da620f0ea
  style ba74afb9_ef25_84a8_66e2_1d25dcee421f fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 2051–2068

    def batch_size_finder(self, device, model_name, initial_batch_size=1024):
        batch_size = initial_batch_size
        while batch_size >= 1:
            empty_gpu_cache(current_device)
            try:
                device, name, model, example_inputs, _ = self.load_model(
                    device,
                    model_name,
                    batch_size,
                )
                self.model_iter_fn(model, example_inputs)
                return batch_size
            except RuntimeError as e:
                error_str = str(e)
                if "channels_last" in error_str:
                    break
            batch_size = self.decay_batch_exp(batch_size)
        return 1