Flux.Conv type instability

[ltjkoomen]

Hi all,

As I was playing around with Flux I noticed that the convolutional layers I was using (Flux.Conv) have some un-inferable types, at least that is my understanding.

The following is a minimal example:

using Flux
imgbatch = rand(Float32, 50, 50, 3, 64)
convlayer = Flux.Conv((3,3), 3=>32)
@code_warntype convlayer(imgbatch)

On my laptop, with Julia version 1.4.1 and Flux version 0.10.4 this results in:

Variables
  c::Conv{2,4,typeof(identity),Array{Float32,4},Array{Float32,1}}
  x::Array{Float32,4}
  #105::Flux.var"#105#106"
  σ::typeof(identity)
  b::Array{Float32,4}
  cdims::DenseConvDims{2,_A,_B,_C,_D,_E,_F,_G} where _G where _F where _E where _D where _C where _B where _A

Body::Any
1 ─ %1  = Base.getproperty(c, :σ)::Core.Compiler.Const(identity, false)
│   %2  = Base.getproperty(c, :bias)::Array{Float32,1}
│   %3  = Core.tuple(%2)::Tuple{Array{Float32,1}}
│         (#105 = %new(Flux.:(var"#105#106")))
│   %5  = #105::Core.Compiler.Const(Flux.var"#105#106"(), false)
│   %6  = Base.getproperty(c, :stride)::Tuple{Int64,Int64}
│   %7  = Flux.map(%5, %6)::Core.Compiler.Const((1, 1), false)
│   %8  = Core.tuple(Flux.:(:), 1)::Core.Compiler.Const((Colon(), 1), false)
│   %9  = Core._apply_iterate(Base.iterate, Flux.reshape, %3, %7, %8)::Array{Float32,4}
│         (σ = %1)
│         (b = %9)
│   %12 = (:stride, :padding, :dilation)::Core.Compiler.Const((:stride, :padding, :dilation), false)
│   %13 = Core.apply_type(Core.NamedTuple, %12)::Core.Compiler.Const(NamedTuple{(:stride, :padding, :dilation),T} where T<:Tuple, false)
│   %14 = Base.getproperty(c, :stride)::Tuple{Int64,Int64}
│   %15 = Base.getproperty(c, :pad)::NTuple{4,Int64}
│   %16 = Base.getproperty(c, :dilation)::Tuple{Int64,Int64}
│   %17 = Core.tuple(%14, %15, %16)::Tuple{Tuple{Int64,Int64},NTuple{4,Int64},Tuple{Int64,Int64}}
│   %18 = (%13)(%17)::NamedTuple{(:stride, :padding, :dilation),Tuple{Tuple{Int64,Int64},NTuple{4,Int64},Tuple{Int64,Int64}}}
│   %19 = Core.kwfunc(Flux.DenseConvDims)::Core.Compiler.Const(Core.var"#Type##kw"(), false)
│   %20 = Base.getproperty(c, :weight)::Array{Float32,4}
│         (cdims = (%19)(%18, Flux.DenseConvDims, x, %20))
│   %22 = σ::Core.Compiler.Const(identity, false)
│   %23 = Base.getproperty(c, :weight)::Array{Float32,4}
│   %24 = Flux.conv(x, %23, cdims)::AbstractArray{yT,4} where yT
│   %25 = Base.broadcasted(Flux.:+, %24, b)::Any
│   %26 = Base.broadcasted(%22, %25)::Any
│   %27 = Base.materialize(%26)::Any
└──       return %27

It seems to be caused by the un-inferable type of cdims, causing the output type of Flux.conv to be inferred as AbstractArray{yT,4} where yT.

While looking at the type definition and constructors of DenseConvDims I noticed that its parametric type is dependent on the contents of, for instance, the variable corresponding to the pad keyword of the Flux.Conv constructor. It might be a stupid question, but why are all the parameters of the type DenseConvDims stored in its type directly instead of in its fields? Doesn't that make the type uncertain by default since the compiler has no way of knowing the contents of variables beforehand? I have only recently started programming in Julia, so please correct me if I am wrong.

[jondeuce]

Bumping this issue, which is still persisting on Julia v1.5.3 and v1.6.0-beta1:

julia-1.6.0-beta1>@code_warntype Flux.DenseConvDims((3,1,1,5), (3,1,1,1))
Variables
  #self#::Type{NNlib.DenseConvDims}
  x_size::NTuple{4, Int64}
  w_size::NTuple{4, Int64}

Body::NNlib.DenseConvDims{2, _A, _B, _C, _D, _E, _F, _G} where _G where _F where _E where _D where _C where _B where _A
1 ─ %1 = NNlib.:(var"#DenseConvDims#6")(1, 0, 1, false, #self#, x_size, w_size)::NNlib.DenseConvDims{2, _A, _B, _C, _D, _E, _F, _G} where _G where _F where _E where _D where _C where _B where _A
└──      return %1

[pxl-th]

Can someone maybe explain, why we need to specialize ConvDims and its derivatives on the kernel size, channels, etc?
Specialization would make sense, IMO, if we are using StaticArray or something that encodes its size in the type. But if we are using regular Array or CuArray then it is not clear why the specialization is needed.
Unless I'm missing something...

Looking, for example, at how the output array in NNlib.conv is created:

y = similar(x, promote_type(xT, wT), output_size(cdims)..., channels_out(cdims), size(x,N))

[DhairyaLGandhi]

These configuration settings can then be specialised on at compile time in many cases, which is why we need these.

[ToucheSir]

Worth looking at which type params are absolutely needed for performance and are used in the wild. For the rest, if one can demonstrate that removing them helps significantly with compilation latency while not compromising the targets in the original PR, that should be a shoe-in.

[pxl-th]

As a proof of concept, I've implemented custom DenseConvDims structure that only encodes spatial_dims in its type, everything else is just a struct field. But it is fully compatible in terms of API with NNlib.conv and others and can be used as a drop-in replacement in Flux for the convolutional layer.

But it allows for full type inference on the Conv layer.
As you can see below, forward pass is greatly improved, while backward remains almost the same, which suggests that the issue is largely in Zygote.

Side note: pooling layers have the same issue with type-inference.

Function used to test:

function main()
    x = zeros(Float32, 28, 28, 1, 5)
    y = zeros(Float32, 10, 5)
    m = Chain(
        Conv((3, 3), 1 => 2, relu),
        Conv((3, 3), 2 => 2, relu),
        Conv((3, 3), 2 => 3, relu),
        Conv((3, 3), 3 => 3, relu),
        Conv((3, 3), 3 => 3, relu),
        Conv((3, 3), 3 => 3, relu),
        x -> reshape(x, :, size(x, 4)),
        Dense(768, 10), softmax)
    θ = params(m)

    @time m(x)
    @time gradient(θ) do
        Flux.crossentropy(m(x), y)
    end
end

Before:

(F) 7.417844 seconds (31.90 M allocations: 1.663 GiB, 10.85% gc time, 99.95% compilation time)
(B) 32.919539 seconds (70.26 M allocations: 3.644 GiB, 4.29% gc time, 99.96% compilation time)

julia> @code_warntype m(x)
MethodInstance for (::Chain{Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}})(::Array{Float32, 4})
  from (c::Chain)(x) in Flux at /home/pxl-th/.julia/dev/Flux/src/layers/basic.jl:49
Arguments
  c::Chain{Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}}
  x::Array{Float32, 4}
Body::Any
1 ─ %1 = Base.getproperty(c, :layers)::Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}
│   %2 = Flux.Tuple(%1)::Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}
│   %3 = Flux.applychain(%2, x)::Any
└──      return %3

After:

(F) 2.192996 seconds (9.31 M allocations: 479.607 MiB, 7.89% gc time, 99.92% compilation time)
(B) 30.723445 seconds (73.09 M allocations: 3.798 GiB, 4.64% gc time, 99.96% compilation time)

julia> @code_warntype m(x)
MethodInstance for (::Chain{Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}})(::Array{Float32, 4})
  from (c::Chain)(x) in Flux at /home/pxl-th/.julia/dev/Flux/src/layers/basic.jl:49
Arguments
  c::Chain{Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}}
  x::Array{Float32, 4}
Body::Matrix{Float32}
1 ─ %1 = Base.getproperty(c, :layers)::Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}
│   %2 = Flux.Tuple(%1)::Tuple{Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Conv{2, 4, typeof(relu), Array{Float32, 4}, Vector{Float32}}, Flux.var"#312#314", Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(softmax)}
│   %3 = Flux.applychain(%2, x)::Matrix{Float32}
└──      return %3

[ToucheSir]

Very nice, IIUC this can be done as a backwards compatible change to NNlib?

For the backwards pass, I guess we'd have to break down how much time is doing into pullback function generation vs pullback execution. At least on my machine, it takes at least 10s just to compile the Zygote compiler itself (for any function).

[pxl-th]

Very nice, IIUC this can be done as a backwards compatible change to NNlib?

I'd say it is breaking if someone expects DenseConvDims to still have these data in its type.
If looking at backwards compatibility from the methods POV, then yes, it is compatible.
I'll open a PR in NNlib as soon as I finish refactoring.