runBenchmark<I, O> function

void runBenchmark<I, O>(String name, Layer<I, O> layer, Tensor<I> input, int iterations,)

Implementation

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

  layer.build(input);

  for (int i = 0; i < 5; i = i + 1) {
    Tensor<O> out = layer.forward(input);
    Tensor<Scalar> loss = pseudoLoss(out);
    loss.backward();
  }

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

  for (int i = 0; i < iterations; i = i + 1) {
    for (int p = 0; p < layer.parameters.length; p = p + 1) {
      layer.parameters[p].zeroGrad();
    }
    input.zeroGrad();

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

    Tensor<Scalar> loss = pseudoLoss(out);

    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');
}