performance_tuning

Quick performance benchmark for GFlowNet training.

Runs per-iteration timing across environments (HyperGrid, DiffusionSampling, DiscreteEBM), losses (TB, SubTB, ModifiedDB), debug on/off, and torch.compile on/off (compiling the standard sampling loop when possible). Produces a bar plot saved to ~/performance_tuning.png.

Attributes

LR

LR_LOGF

LR_LOGZ

Classes

BenchmarkConfig

BenchmarkResult

BenchmarkSettings

Functions

build_components(env_name, loss_name, debug, device)

Create env, gflownet, optimizer for a given setup.

main()

make_step_fn(env, gflownet, optimizer, batch_size, device)

maybe_compile(fn)

plot_results(results[, output_path, return_fig])

resolve_device([requested])

run_benchmark(config, device, settings)

run_benchmarks(settings, *[, return_fig, verbose])

synchronize(device)

time_iterations(step_fn, device, warmup, n_iters)

Module Contents

class performance_tuning.BenchmarkConfig
debug: bool
env_name: str
property label: str
Return type:

str

loss_name: str
use_compile: bool
class performance_tuning.BenchmarkResult
compiled: bool
config: BenchmarkConfig
elapsed: float
env_name: str
mean_iter_ms: float
reason: str | None = None
skipped: bool
std_iter_ms: float
class performance_tuning.BenchmarkSettings
batch_size: int = 32
deterministic_mode: bool = False
device: str = 'cpu'
n_iters: int = 200
output_path: pathlib.Path | None
warmup_iters: int = 50
performance_tuning.LR = 0.001
performance_tuning.LR_LOGF = 0.001
performance_tuning.LR_LOGZ = 0.1
performance_tuning.build_components(env_name, loss_name, debug, device)

Create env, gflownet, optimizer for a given setup.

Parameters:

  • env_name (str)

  • loss_name (str)

  • debug (bool)

  • device (torch.device)

performance_tuning.main()
Return type:

None

performance_tuning.make_step_fn(env, gflownet, optimizer, batch_size, device)
Parameters:

  • batch_size (int)

  • device (torch.device)

Return type:

Callable[[], float]

performance_tuning.maybe_compile(fn)
Parameters:

fn (Callable[[], float])

Return type:

tuple[Callable[[], float], bool]

performance_tuning.plot_results(results, output_path=None, return_fig=False)
Parameters:

  • results (list[BenchmarkResult])

  • output_path (pathlib.Path | None)

  • return_fig (bool)

performance_tuning.resolve_device(requested='auto')
Parameters:

requested (str)

Return type:

torch.device

performance_tuning.run_benchmark(config, device, settings)
Parameters:
Return type:

BenchmarkResult

performance_tuning.run_benchmarks(settings, *, return_fig=False, verbose=False)
Parameters:
performance_tuning.synchronize(device)
Parameters:

device (torch.device)

Return type:

None

performance_tuning.time_iterations(step_fn, device, warmup, n_iters)
Parameters:

  • step_fn (Callable[[], float])

  • device (torch.device)

  • warmup (int)

  • n_iters (int)

Return type:

tuple[float, list[float]]