diff options
Diffstat (limited to 'test/benchmarks/ml')
| -rw-r--r-- | test/benchmarks/ml/tensorflow_test.go | 15 |
1 files changed, 8 insertions, 7 deletions
diff --git a/test/benchmarks/ml/tensorflow_test.go b/test/benchmarks/ml/tensorflow_test.go index a55329d82..b0e0c4720 100644 --- a/test/benchmarks/ml/tensorflow_test.go +++ b/test/benchmarks/ml/tensorflow_test.go @@ -22,8 +22,6 @@ import ( "gvisor.dev/gvisor/test/benchmarks/harness" ) -var h harness.Harness - // BenchmarkTensorflow runs workloads from a TensorFlow tutorial. // See: https://github.com/aymericdamien/TensorFlow-Examples func BenchmarkTensorflow(b *testing.B) { @@ -38,7 +36,7 @@ func BenchmarkTensorflow(b *testing.B) { "NeuralNetwork": "3_NeuralNetworks/neural_network.py", } - machine, err := h.GetMachine() + machine, err := harness.GetMachine() if err != nil { b.Fatalf("failed to get machine: %v", err) } @@ -49,15 +47,17 @@ func BenchmarkTensorflow(b *testing.B) { ctx := context.Background() b.ResetTimer() + b.StopTimer() + for i := 0; i < b.N; i++ { - b.StopTimer() container := machine.GetContainer(ctx, b) defer container.CleanUp(ctx) if err := harness.DropCaches(machine); err != nil { b.Skipf("failed to drop caches: %v. You probably need root.", err) } - b.StartTimer() + // Run tensorflow. + b.StartTimer() if out, err := container.Run(ctx, dockerutil.RunOpts{ Image: "benchmarks/tensorflow", Env: []string{"PYTHONPATH=$PYTHONPATH:/TensorFlow-Examples/examples"}, @@ -65,13 +65,14 @@ func BenchmarkTensorflow(b *testing.B) { }, "python", workload); err != nil { b.Fatalf("failed to run container: %v logs: %s", err, out) } + b.StopTimer() } }) } - } func TestMain(m *testing.M) { - h.Init() + harness.Init() + harness.SetFixedBenchmarks() os.Exit(m.Run()) } |