# Flux: Locally connected layer performance

**URL:** <https://discourse.julialang.org/t/flux-locally-connected-layer-performance/55711>\
**Category:** Machine Learning\
**Tags:** gpu, performance, flux, tullio, loopvectorization\
**Created:** [February 21, 2021, 11:25am UTC](https://discourse.julialang.org/t/flux-locally-connected-layer-performance/55711 "2021-02-21T11:25:40Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![johnbb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/johnbb/32/34233_2.png) [@johnbb](https://discourse.julialang.org/u/johnbb)\
**Post date:** [February 21, 2021, 11:25am UTC](https://discourse.julialang.org/t/flux-locally-connected-layer-performance/55711/1 "2021-02-21T11:25:40Z")

</div>

I have tried to implement a 2D locally connected layer (convolution with non-shared weights) in Flux restricted to filter size (1,1) only. That is, the linear combinations/convolutions are only over the (3rd) channel/variable dimension. Since this type of layer involves many weight parameters, an efficient implementation with respect to speed and memory management is essential. Is there any scope for improvement of the following code, in particular for GPU?

```julia
function (a::Local2D)(x::AbstractArray) # size(x) = (w,h,cin,n), size(W) = (w,h,cin,cout)
    W, b, σ = a.W, a.b, a.σ
    out = cat([sum(W[:,:,:,i] .* x, dims=3) .+ b[:,:,:,i] for i in 1:size(W,4)]..., dims=3)
    return σ.(out)
end

```

The complete code including a test example:

```julia
struct Local2D{S<:AbstractArray, T<:AbstractArray, F}
  W::S
  b::T
  σ::F
end

Local2D(W, b) = Local2D(W, b, identity)

function Local2D(w::Integer, h::Integer, cin::Integer, cout::Integer, σ = identity;
                 initW = Flux.glorot_uniform, initb = zeros)
    return Local2D(initW(w,h,cin,cout), initb(Float32, w,h,1,cout), σ)
end

Flux.@functor Local2D

function (a::Local2D)(x::AbstractArray) # size(x) = (w,h,cin,n), size(W) = (w,h,cin,cout)
    W, b, σ = a.W, a.b, a.σ
    out = cat([sum(W[:,:,:,i] .* x, dims=3) .+ b[:,:,:,i] for i in 1:size(W,4)]..., dims=3)
    return σ.(out)
end

# example
device = gpu
n = 1024; d = 256; cin = 4; cout = 2
x = rand(Float32, d, d, cin, n) |> device;
y = rand(Float32, d, d, cout, n) |> device;
trdata = Flux.Data.DataLoader((x, y), batchsize=32) |> device;
model = Local2D(d, d, cin, cout, identity) |> device;
loss(x, y) = Flux.mse(model(x), y)
@time Flux.train!(loss, Flux.params(model), trdata, ADAM())

```

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [February 21, 2021, 12:08pm UTC](https://discourse.julialang.org/t/flux-locally-connected-layer-performance/55711/2 "2021-02-21T12:08:43Z")

</div>

It should be possible to do this without making slices, and this is likely to be quicker. A few ways are:

```julia
julia> x = rand(256,256,4,32); W = rand(256,256,4,2);

julia> using TensorCast, OMEinsum, Tullio, LoopVectorization

julia> @reduce mid_1[i,j,k,n] := sum(c) W[i,j,c,k] * x[i,j,c,n];

julia> @ein mid_2[i,j,k,n] := W[i,j,c,k] * x[i,j,c,n];

julia> @tullio mid_3[i,j,k,n] := W[i,j,c,k] * x[i,j,c,n];

julia> size(mid_1)
(256, 256, 2, 32)

julia> mid_1 ≈ mid_2 ≈ mid_3 ≈ cat([sum(W[:,:,:,i] .* x, dims=3) for i in 1:size(W,4)]..., dims=3)
true

julia> out = σ.(mid_1 .+ b) ^C

```

These do different things, `@reduce` is just broadcasting, it first uses `reshape(x, size(x)[1:3]...,1,:)` to place its 4th dimension 5th. `@ein` will (I think) call batched matrix multiplication here, after using `permutedims` to line things up. (You could write either of those out by hand, too.) `@tullio` generates loops directly.

```julia
julia> @pretty @reduce mid_1[i,j,k,n] := sum(c) W[i,j,c,k] * x[i,j,c,n];
begin
    local lyrebird = orient(x, (:, :, :, *, :))
    mid_1 = dropdims(sum(@ __dot__ (W * lyrebird), dims = 3), dims = 3)
end

julia> @btime @reduce mid_1[i,j,k,n] := sum(c) $W[i,j,c,k] * $x[i,j,c,n];
  70.918 ms (42 allocations: 160.00 MiB)

julia> @btime @ein mid_2[i,j,k,n] := $W[i,j,c,k] * $x[i,j,c,n];
  36.443 ms (324 allocations: 132.03 MiB)

julia> @btime @tullio mid_3[i,j,k,n] := $W[i,j,c,k] * $x[i,j,c,n];
  10.172 ms (76 allocations: 32.01 MiB)

julia> @btime cat([sum($W[:,:,:,i] .* $x, dims=3) for i in 1:size($W,4)]..., dims=3);
  56.664 ms (79 allocations: 196.00 MiB)

```

How fast these are will depend a lot on what size the arrays are, esp. how many channels. Times on the GPU may also look very different, and the gradient will be slower to a varying degree.

---

<div class="post-metadata">

**Author:** ![johnbb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/johnbb/32/34233_2.png) [@johnbb](https://discourse.julialang.org/u/johnbb)\
**Post date:** [February 21, 2021, 7:34pm UTC](https://discourse.julialang.org/t/flux-locally-connected-layer-performance/55711/3 "2021-02-21T19:34:20Z")

</div>

Thanks @mcabbott! A very brief test on my example indicates that `@tullio` seems most promising for CPU and `@ein` for GPU. More extensive tests are needed though.