Flux.jl: Batch aware binarycrossentropy and logitbinarycrossentropy

[sambitdash]

While the crossentropy methods are batch aware the binarycrossentropy and logitbinarycrossentropy work on individual scalars. This may not be a good idea to do so from an interface standpoint.

I feel the logitbinarycrossentropy and binarycrossentropy must be made batch aware as well.

[CarloLucibello]

Yep, we should be more consistent. Ideally, for each loss we should have a reduce keyword to specify the type of reduction, e.g. mean or sum

cf. https://pytorch.org/docs/stable/nn.html#crossentropyloss

[saswatpp]

I would love to work on this issue. Just a small doubt: Are we looking for simple if-else conditions inside functions to check for reduction method(sum/mean/none)? Should we also extend logitbinarycrossentropy and binarycrossentropy for different weight types and make separate Method definitions for AbstractVecOrMat type? @CarloLucibello

[CarloLucibello]

I think most loss functions here need a revamp and unified interface. For instance, we should remove the restriction to VecOrMat and accept arbitrary tensors.

We should have a reduction keyword argument, which can take a function input (typically the functions mean or sum) or nothing. It should default to mean. This way we should only branch out the nothing case. Every loss should have this. Later on, we could also add a keyword argument rdims to specify which dimensions we should reduce over. rdims=: as default.

Further addition for crossentropy: have a dims::Integer=1 keyword specifying which dimension is representing the probabilities

[saswatpp]

Will something like this do? Should I extend it for other functions too ?

function apply_loss_reduction(loss, reduction)
  if reduction=="mean"
    return loss = sum(loss) * 1//length(loss)
  elseif reduction=="sum"
    return loss = sum(loss)
  else
    return loss
  end
end

function mse2(y::AbstractArray, ŷ::AbstractArray; reduction="mean")
  return apply_loss_reduction((ŷ .- y).^2, reduction)
end

@CarloLucibello

[CarloLucibello]

The reduction argument could be a function, e.g. (mean or sum) or nothing for no reduction.
We could also add a rdims keyword to specify reduction dimensions.
For instance in softmax you would have both a dims (for the probability dimension) and a rdims keyword argument with defaults as follows:

function softmax(ŷ, y; dims=1, reduction=mean, rdims=:)
...

We should think carefully about the design here because we want to get it right and avoid changing it later.

[saswatpp]

Will have to add sum and mean functions to utils.jl

#utils.jl
sum(A::AbstractArray; dims=:) = Base.sum(A, dims=dims)
sum(A::AbstractArray, ::Colon) = Base.sum(A)

mean(A::AbstractArray; dims=:) = mean(A, dims)
mean(A::AbstractArray, ::Colon) = sum(A) / length(A)
mean(A::AbstractArray, dims) = sum(A, dims=dims) / mapreduce(i -> size(A, i), *, unique(dims); init=1)

#Stateless.jl
function mse(y::AbstractArray, ŷ::AbstractArray; reduction=mean, rdims=:)
    return reduction(((ŷ .- y).^2), dims=rdims)
end
mse(y::AbstractArray, ŷ::AbstractArray, reduction::Nothing) = ((ŷ .- y).^2)

For mse(a,b,reduction=nothing) it gives Error function Nothing does not accept keyword arguments. Any way to eliminate this. Can extend this to other loss functions if you say. @CarloLucibello

[saswatpp]

Had to extend Core.nothing to avoid the error and reduce Number of methods for loss function to 1(No need to branch Out)

#Utils.jl
sum(A::AbstractArray; dims=:) = Base.sum(A, dims=dims)
sum(A::AbstractArray, ::Colon) = Base.sum(A)

mean(A::AbstractArray; dims=:) = mean(A, dims)
mean(A::AbstractArray, ::Colon) = sum(A) / length(A)
mean(A::AbstractArray, dims) = sum(A, dims=dims) / mapreduce(i -> size(A, i), *, unique(dims); init=1)

import Core:nothing
nothing(A::AbstractArray; dims) = A

#Stateless.jl
function mse(y::AbstractArray, ŷ::AbstractArray; reduction=mean, rdims=:)
    return reduction(((ŷ .- y).^2), dims=rdims)
end

This could be used to extend existing Loss functions. @CarloLucibello

[CarloLucibello]

That is type piracy, you should always avoid that.
Also, we already have mean and sum functions in Base, I don't see why you should define them.
This is one possibility:

function mse(y::AbstractArray, ŷ::AbstractArray; reduction=mean, rdims=:)
   if reduction !== nothing
         reduction(((ŷ .- y).^2), dims=rdims)
   else
         (ŷ .- y).^2
   end
end

ps: you can use ```julia for julia code block fencing.

[CarloLucibello]

sorry, I see now that the above doesn't work, let me think about it

[CarloLucibello]

Actually, it does work (just need a using Statistics).
It is also type-stable (check with @code_warntype)

[saswatpp]

So I tried Implementing this on BinaryCrossEntropy

_binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)
function binarycrossentropy(ŷ::AbstractArray, y::AbstractArray; reduction=mean, rdims=:)
   if reduction !== nothing
        reduction(_binarycrossentropy.(ŷ, y), dims=rdims)
   else
        _binarycrossentropy.(ŷ, y)
   end
end

It tested correctly for following example with Pytorch nn.BCELoss(). Had to remove Keyword arg ϵ just like in Pytorch version where it is not required.

binarycrossentropy(σ.(Float32[-1.1491, 0.8619, 0.3127]),Float32[1, 1, 0.],reduction=nothing)

3-element Array{Float32,1}:
 1.424397  
 0.35231656
 0.8616702 

[sambitdash]

@saswatpp can the reduction be made identity function than nothing. In that case, special check for nothing may not be needed.

Any reasons, why $ \epsilon $ was removed? Is it just because it's not there in PyTorch does not seem reason enough. It's never a good idea to remove parameters from an established API as that may break people's code who are already using the API.

[saswatpp]

identity function reduces understandability. Also Keyword args are not supported in identity function so it won't even work in the first place. Extended it is out of question.

Now for the ϵ, we have to decide on default value for ϵ if we want to put it in error function. As that single ϵ value will be used for every elementwise Calculation.

[DhairyaLGandhi]

Changing the API is breaking, and identity is a well understood function, so it should be fine and just be a no op. It not supporting kwargs should just be designed such that kwargs are empty when splatted to it which should just be compiled away I think.

[CarloLucibello]

I favor nothing over identity: reduction=nothing makes the point that no reduction is applied, while reduction=identity it's not semantically consistent with reduction="something lowering the dimensionality of the input" as reductions are commonly understood. I would still be ok with identity though, just prefer nothing.

As for breaking changes, we should definitely break binarycrossentropy and logitbinarycrossentropy, since they are inconsistent with the other losses performing mean reduction. The sooner we do it the better

[saswatpp]

mse, mae, Binarycrossentropy, logitbinarycrossentropy, hinge, squaredhinge and poisson losses seem to be modified to have reduction and rdims kwarg. Should I compile them into a PR ? @CarloLucibello

[saswatpp]

Also need to start working on Multidimensional Cross entropy loss if you could give suggestions.

[CarloLucibello]

closed in #1150