# Is there a way to speed up this implementation of BLSTM layer?

**URL:** <https://discourse.julialang.org/t/is-there-a-way-to-speed-up-this-implementation-of-blstm-layer/36429>\
**Category:** Machine Learning\
**Created:** [March 24, 2020, 10:36am UTC](https://discourse.julialang.org/t/is-there-a-way-to-speed-up-this-implementation-of-blstm-layer/36429 "2020-03-24T10:36:45Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Azamat](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/azamat/32/6892_2.png) [@Azamat](https://discourse.julialang.org/u/Azamat)\
**Post date:** [March 24, 2020, 10:36am UTC](https://discourse.julialang.org/t/is-there-a-way-to-speed-up-this-implementation-of-blstm-layer/36429/1 "2020-03-24T10:36:45Z")

</div>

I have implemented the BLSTM layer for a batch of sequences input stored as a 3D array with dimensions ordered as `[features × time × batch]`. And I wonder if there is a faster way to implement it. I’m interested in the speed of both forward and backward passes. Here is my implementation:

```nohighlight
using Flux
using Flux: @functor, Recur, LSTMCell
using Zygote
using Zygote: Buffer

struct BLSTM{M <: DenseMatrix, V <: DenseVector}
   forward :: Recur{LSTMCell{M,V}}
   backward :: Recur{LSTMCell{M,V}}
   outdim :: Int
end

Flux.trainable(m::BLSTM) = (m.forward, m.backward)
@functor BLSTM

function BLSTM(in::Integer, out::Integer)
   forward = LSTM(in, out)
   backward = LSTM(in, out)
   return BLSTM(forward, backward, Int(out)) |> gpu
end

function (m::BLSTM)(Xs::DenseArray{<:Real,3})
   Xs = permutedims(Xs, (1, 3, 2)) # [features × time × batch] -> [features × batch × time]
   # preallocate output buffer
   Ys = Buffer(Xs, 2m.outdim, size(Xs,3), size(Xs,2))
   axisYs₁ = axes(Ys, 1)
   time_f = axes(Ys, 2)
   time_b = reverse(time_f)
   @inbounds begin
      # get forward and backward slice indices
      slice_f = axisYs₁[1:m.outdim]
      slice_b = axisYs₁[(m.outdim+1):end]
      # bidirectional run step
      setindex!.((Ys,), m.forward.(view.((Xs,), :, :, time_f)), (slice_f,), time_f, :)
      setindex!.((Ys,), m.backward.(view.((Xs,), :, :, time_b)), (slice_b,), time_b, :)
      # the same as
      # @views for (t_f, t_b) ∈ zip(time_f, time_b)
      # Ys[slice_f, t_f, :] = m.forward(Xs[:, :, t_f])
      # Ys[slice_b, t_b, :] = m.backward(Xs[:, :, t_b])
      # end
      # but implemented via broadcasting as Zygote differentiates loops much slower than broadcasting
   end
   return copy(Ys) # [features × time × batch]
end

```

And here are the results I get:

```nohighlight
using BenchmarkTools

D, T, B = 512, 800, 32
m = BLSTM(D, D÷2)
θ = params(m)
Xs = rand(Float32, 512, 800, 32) |> gpu

f(m, Xs) = sum(m(Xs))
g(m, Xs, θ) = gradient(() -> f(m, Xs), θ)

julia> @benchmark f($m, $Xs)
BenchmarkTools.Trial: 
  memory estimate: 26.50 MiB
  allocs estimate: 722527
  --------------
  minimum time: 607.991 ms (0.00% GC)
  median time: 629.910 ms (1.41% GC)
  mean time: 625.312 ms (1.13% GC)
  maximum time: 631.088 ms (1.46% GC)
  --------------
  samples: 9
  evals/sample: 1

julia> Flux.reset!(m); @benchmark g($m, $Xs, $θ)
BenchmarkTools.Trial: 
  memory estimate: 28.98 MiB
  allocs estimate: 797407
  --------------
  minimum time: 603.341 ms (0.00% GC)
  median time: 622.730 ms (1.30% GC)
  mean time: 621.413 ms (1.17% GC)
  maximum time: 629.907 ms (1.38% GC)
  --------------
  samples: 9
  evals/sample: 1

```

Any improvement suggestions?

---

<div class="post-metadata">

**Author:** ![Azamat](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/azamat/32/6892_2.png) [@Azamat](https://discourse.julialang.org/u/Azamat)\
**Post date:** [March 28, 2020, 11:00am UTC](https://discourse.julialang.org/t/is-there-a-way-to-speed-up-this-implementation-of-blstm-layer/36429/2 "2020-03-28T11:00:46Z")

</div>

I discovered that the speed up above is due to bug in Zygote, which is dropping the gradients. Filed [here](https://github.com/FluxML/Zygote.jl/issues/563).