runBenchmark<I, O> function

void runBenchmark<I, O>(String name, Layer<I, O> layer, Tensor<I> input, int iterations, { bool showGraph = false, })

Implementation

void runBenchmark<I, O>(String name, Layer<I, O> layer, Tensor<I> input, int iterations, {bool showGraph = false}) {
  print('--- Benchmarking: $name ---');

  // 1. Build the layer
  layer.build(input);

  // 2. Warmup (to trigger JIT compilation and memory allocation)
  for (int i = 0; i < 5; i = i + 1) {
    Tensor<O> out = layer.forward(input);
    Tensor<Scalar> loss = pseudoLoss(out);

    // Print the compute graph on the very first warmup pass if requested
    if (showGraph && i == 0) {
      loss.printGraph();
      print('');
    }

    loss.backward();
  }

  Stopwatch fwWatch = Stopwatch();
  Stopwatch bwWatch = Stopwatch();

  // 3. Main Benchmark Loop
  for (int i = 0; i < iterations; i = i + 1) {
    // Zero gradients before each step
    for (int p = 0; p < layer.parameters.length; p = p + 1) {
      layer.parameters[p].zeroGrad();
    }
    input.zeroGrad();

    // Time Forward Pass
    fwWatch.start();
    Tensor<O> out = layer.forward(input);
    fwWatch.stop();

    Tensor<Scalar> loss = pseudoLoss(out);

    // Time Backward Pass
    bwWatch.start();
    loss.backward();
    bwWatch.stop();
  }

  double avgFw = fwWatch.elapsedMilliseconds / iterations;
  double avgBw = bwWatch.elapsedMilliseconds / iterations;

  print('Forward Pass:  ${avgFw.toStringAsFixed(2)} ms / step');
  print('Backward Pass: ${avgBw.toStringAsFixed(2)} ms / step');
  print('Total Time:    ${(avgFw + avgBw).toStringAsFixed(2)} ms / step\n');
}