Recurrent benchmarks

perf/bench_utils.jl

@@ -1,12 +1,24 @@
 using BenchmarkTools
 using Flux
 using CUDA
-using Zygote: pullback
+using Zygote: pullback, ignore
 
 
 fw(m, x) = m(x)
 bw(back) = back(1f0)
 fwbw(m, ps, x) = gradient(() -> sum(m(x)), ps)
+
+# Need to specialize for flux.recur.
+fw(m::Flux.Recur, X::Vector{<:AbstractArray}) = begin
+    ignore() do
+      Flux.reset!(m)
+    end
+    [m(x) for x in X]
+end
+fwbw(m::Flux.Recur, ps, X::Vector{<:AbstractArray}) = gradient(ps) do
+    y = fw(m, X)
+    sum(sum(y))
+end
 
 function run_benchmark(model, x; cuda=true)
 
@@ -16,7 +28,11 @@ function run_benchmark(model, x; cuda=true)
     end
 
     ps = Flux.params(model)
-    y, back = pullback(() -> sum(model(x)), ps)
+    y, back = if model isa Flux.Recur && eltype(x) <: AbstractArray
+        pullback(() -> sum(sum([model(x_t) for x_t in x])), ps)
+    else
+        pullback(() -> sum(model(x)), ps)
+    end
 
 
     if cuda

perf/recurrent.jl

@@ -0,0 +1,41 @@
+
+println("RNN")
+for n in [2, 20, 200, 1000], T in [1, 8, 16, 64]
+  x = [randn(Float32, n, n) for t in 1:T]
+  model = RNN(n, n)
+  println("CPU n=$n, t=$T")
+  run_benchmark(model, x, cuda=false)
+  println("CUDA n=$n, t=$T")
+  try
+      run_benchmark(model, x, cuda=true)
+  catch ex
+      @show typeof(ex)
+      if ex isa OutOfGPUMemoryError
+          @warn "Not enough GPU memory to run test"
+      else
+          rethrow(ex)
+      end
+  end
+end
+
+println("RNN-3d")
+for n in [2, 20, 200, 1000], T in [1, 8, 16, 64]
+  x = randn(Float32, n, n, T)
+  model = RNN(n, n)
+  println("CPU n=$n, t=$T")
+  run_benchmark(model, x, cuda=false)
+  println("CUDA n=$n, t=$T")
+  try
+      run_benchmark(model, x, cuda=true)
+  catch ex
+      @show typeof(ex)
+      if ex isa OutOfGPUMemoryError
+          @warn "Not enough GPU memory to run test"
+      else
+          rethrow(ex)
+      end
+  end
+end
+
+
+

perf/runbenchmarks.jl

@@ -11,3 +11,6 @@ include("conv.jl")
 
 @info "Benchmark VGG"
 include("vgg.jl")
+
+@info "Benchmark Recurrent"
+include("recurrent.jl")