outdims function doesn't work properly for chained layers

[HamletWantToCode]

In PR #960 outdims function is added to Flux.jl and makes it easier for us to infer the output dimension for our neural network model ( especially when using convolutional layer ). However, the results of outdims function is incorrect for chained layers, and existing test cases fail to reflect this flaw. e.g. Try following code:

using Flux
using Flux: outdims

all_dense_model = Chain(Dense(2, 10), Dense(10, 4))
outdims(all_dense_model, 2)    # will return (10,) while it should be (4,)

I'm using Flux v0.10.3 in Julia v1.3.0 under macOS

Currently, outdims is implemented as:

outdims(c::Chain, isize) = foldl(∘, map(l -> (x -> outdims(l, x)), c.layers))(isize)

however, foldl doesn't give correct function composition, e.g. Try following code:

f1(x) = 2x; f2(x) = exp(x); f3(x) = x^3
const x0 = 0.5
# we want to compute f3∘(f2∘f1(x))
f_true = f3∘f2∘f1
f_true(x0)   # 20.085536923187664
# however
f = foldl(∘, (f1, f2, f3))   # the returned f = ((f1∘f2)∘f3) following left association rule !!!
f(x0) # 2.2662969061336526

I suggest to resolve this problem by the following modifications:

# edit
using Base: tail # thanks to @darsnack for pointing this out

outdims(t::Tuple, isize) = outdims(tail(t), outdims(first(t), isize))
outdims(c::Chain, isize) = outdims(c.layers, isize)

and add on following test cases:

@testset "dense layer" begin
         X = randn(Float32, 5, 10)
         D0, D1, D2 = 5, 100, 25
         dense1 = Dense(D0, D1, relu)
         dense2 = Dense(D1, D2)
         dense_chain = Chain(dense1, dense2)
         @test outdims(dense1, D0) = (D1,)
         @test first(outdims(dense1, D0)) == size(dense1(X), 1)
         @test first(outdims(dense_chain, D0)) == size(dense_chain(X), 1)
end

@testset "conv layer" begin
            X = randn(Float32, 28, 28, 1, 1)
            D0, D1 = 3, 5
            S, P = 3, 1
            conv1_stride = Conv((D0, D0), 16=>32, stride=S, pad=P)
            conv2 = Conv((D1, D1), 3=>16)
            conv_chain = Chain(conv1_stride, conv2)
            @test typeof(outdims(conv1_stride, (28, 28))) <: Tuple
            @test length(outdims(conv1_stride, (28, 28))) == 2
            @test outdims(conv1_stride, (28, 28)) == size(conv1_stride(X))[1:2]
            @test outdims(conv_chain, (28, 28)) == size(conv_chain(X))[1:2]
end

Maybe we could also consider extend outdims function to more complex chained model, e.g. model that contain normal julia function, such as x->reshape(x, :, 4) , and also write outdims for recurrent layers.

I'll collect these into a PR if you'd like to @MikeInnes @baggepinnen @darsnack .

[DhairyaLGandhi]

Sounds like a good idea to collect all the failing test cases, having something like outdims(::type of(reshape), ...) might be fine too

[mcabbott]

Can it be taught to understand x->reshape(x, :, 4)? Perhaps people would have to write Base.Fix2(reshape, (:,4)). Right now Flux.outdims(x -> reshape(x,:,4), 5) is an error, which is good, but perhaps a more helpful one would be better.

BTW why is Flux.outdims(Dense(2,5), (3,10)) == (5,)? I would have expect a 2-tuple, and probably an error!

[HamletWantToCode]

Hi @mcabbott

why is Flux.outdims(Dense(2,5), (3,10)) == (5,)

Current implementation of outdims for Dense layer ( see below ) doesn't really depend on the 2nd argument you passed in, it will just ignore (3, 10) and return size(l.W)[1].

outdims(l::Dense, isize) = (size(l.W)[1],)

I think this can be solved by

function outdims(l::Dense, isize::Tuple)
    isize == (size(l.W, 2),) || throw(DimensionMismatch("input size should equal to $((size(l.W)[2],)), got $isize"))
    return (size(l.W, 1),)
end

I prefer to throw an error instead of return a 2-tuple.

[mcabbott]

OK, so outdims does not propagate the size of tensors, it only cares about the first one or two dimensions, depending on the layer. Except that sometimes it checks the others?

julia> Conv((3, 3), 3 => 16)(ones(10, 10, 3, 1)) |> size
(8, 8, 16, 1)

julia> Flux.outdims(Conv((3, 3), 3 => 16), (10, 10, 3, 1))
(8, 8)

julia> Flux.outdims(Conv((3, 3), 3 => 16), (10, 10)) # so it silently ignores 3rd & 4th?
(8, 8)

julia> Flux.outdims(Conv((3, 3), 3 => 16), (10, 10, 10, 1)) # no it doesn't!
ERROR: DimensionMismatch("Input channels must match! (10 vs. 3)")

julia> Flux.outdims(Dense(8, 16), (8,8,16,1)) # but here it ignores everything...
(16,)

Perhaps this is orthogonal to your question, sorry, but at very least I think the docstring should explain what this function is for, "output dimensions" is pretty vague. Why it isn't just outsize(f, s::Tuple) = size(f(ones(Float32, s...))), or faster versions of that? Then x->reshape(x, :, 4) would just work, whereas at present the output seems to depend on things that the layer before will have dropped.

[HamletWantToCode]

Hi @mcabbott

it only cares about the first one or two dimensions, depending on the layer. Except that sometimes it checks the others?

I also think so.

Why it isn't just outsize(f, s::Tuple) = size(f(ones(Float32, s...)))

For the whole network, I think this may results in large overhead, since every time one need to run the complete computation for it's output dimension. However, I think this may be a good idea for functions like x->reshape(x, :, 4), we can just define something like:

outdims(f::typeof(reshape), s::Tuple) = size(f(ones(Float32, s...)))

since this won't cost much.

[DhairyaLGandhi]

Or just return the shape the reshape is wants to reshape it to? Accounting for any Colons, of course

[mcabbott]

I don't think x->reshape(x, :, 4) has a type you can recognise, hence my Base.Fix2 comment above.

What I meant about outsize was that this would be a simple rule for what the function should return, and a fallback implementation, even if it gets skipped on known types such as Conv. But this rule is quite different to what outdims presently returns. Its action on Conv means that it cannot distinguish these two cases:

julia> m = Chain(Conv((3, 3), 3 => 16), Base.Fix2(reshape, (:,4)));

julia> m(ones(10,10,3,1)) |> size
(256, 4)

julia> m(ones(10,10,3,2)) |> size
(512, 4)

[DhairyaLGandhi]

Yes, you are right. The more I think about it, the more I feel for this to work generically and correctly, it would need to do a true forward pass, which sounds like too much ado for something one expects to be trivial and quick. Writing all these specific rules feels incorrect too and won't scale.

One could always have something that does a true forwards pass only as a fallback and only return the size of the output to the next layer, but that feels inelegant too.

[MikeInnes]

Yeah, this is an inherent problem with this kind of API. We definitely don't want to go down the route of supporting things like x -> reshape(x, ...) and so on.

Probably best to recommend that people just make an array and throw it at the model to see what happens, which is the only really reliable source of truth.

[mcabbott]

Having a slowish fallback plus fast shortcuts doesn't seem that inelegant to me. But perhaps I'm not so sure what this is for anyway. Is it for interactive use to check you've defined the model how you think you did? If so, is the speed of simply running it once a big deal?

The present design seems pretty confusing though, it does not seem clearly specified what you should provide, or are going to get back.

[DhairyaLGandhi]

I think either the use case for this needs to be something more concrete where a forwards pass is untenable, or it's implantation needs to be correct, simple and reliable. Perhaps good to clear that bit first?

[darsnack]

The original PR was supposed to address programmatic layer building for large NNs. If you are putting together a series of convolutions, then you need to keep track of the (width, height) of each layer's output so that you can size the final dense layers properly. Based on Slack chatter, it was clear people kept defining this in their own code so it made sense to provide a utility function in Flux.

But perhaps I'm not so sure what this is for anyway.

The only practical use case originally envisioned is what I highlighted above. Thus, it only needs to be there for conv-type layers. When doing the original PR, we thought why don't we provide it for other layers too to be consistent. That's why dense, etc. were added.

To address the input-output relationship. The intent is that you put in (width, height) and get out (width, height). Probably it is a better design to return the full tensor size.

As for the errors, we should be returning errors in all dimension mismatch cases. The reason it does for convolutions and not for dense is because the convolution layers utilize the underlying NNlib functions which throw a dimension mismatch.

Probably best to recommend that people just make an array and throw it at the model to see what happens, which is the only really reliable source of truth.

Since the utility of these functions is during model building, we want fast responses.

[darsnack]

Also, based on the utility described above, I would argue that this function should be restricted to NN layers. I would suggest keeping it defined to dense/conv functions only, and defining a fallback method that determines the output size by evaluating the function. And we should be explicit in the docs that this is what happens for all types not listed in the docs. I am not sure evaluating reshape is considered slow compared to evaluating a series of convolution layers.

Of course, we would include the consistency fixes to the API such as returning the full output tensor size and throwing errors.

[mcabbott]

OK, sounds sensible. If it's going to be defined for both Conv and Dense, then it has to have some way of digesting reshape as you must be mixing 4-tensors and 2-tensors.

Perhaps for cases where you'd like to think only about height/width, not all dimensions, it should do something like outdims(Conv((3, 3), 3 => 16), (10, 10)) == (8, 8, missing, missing).