clean losses

Author: CarloLucibello

docs/make.jl

@@ -1,6 +1,11 @@
 using Documenter, Flux, NNlib, Functors, MLUtils
 
 DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive = true)
+DocMeta.setdocmeta!(Flux.Losses, :DocTestSetup, :(using Flux.Losses); recursive = true)
+
+# In the Losses module, doctests which differ in the printed Float32 values won't fail
+DocMeta.setdocmeta!(Flux.Losses, :DocTestFilters, :(r"[0-9\.]+f0"); recursive = true)
+
 makedocs(modules = [Flux, NNlib, Functors, MLUtils],
          doctest = false,
          sitename = "Flux",

docs/src/models/advanced.md

@@ -213,13 +213,14 @@ model(gpu(rand(10)))
 A custom loss function for the multiple outputs may look like this:
 ```julia
 using Statistics
+using Flux.Losses: mse
 
 # assuming model returns the output of a Split
 # x is a single input
 # ys is a tuple of outputs
 function loss(x, ys, model)
   # rms over all the mse
   ŷs = model(x)
-  return sqrt(mean(Flux.mse(y, ŷ) for (y, ŷ) in zip(ys, ŷs)))
+  return sqrt(mean(mse(y, ŷ) for (y, ŷ) in zip(ys, ŷs)))
 end
 ```

docs/src/models/losses.md

@@ -3,8 +3,16 @@
 Flux provides a large number of common loss functions used for training machine learning models.
 They are grouped together in the `Flux.Losses` module.
 
-Loss functions for supervised learning typically expect as inputs a target `y`, and a prediction `ŷ`.
-In Flux's convention, the order of the arguments is the following
+As an example, the crossentropy function for multi-class classification that takes logit predictions (i.e. not [`softmax`](@ref)ed)
+can be imported with
+
+```julia
+using Flux.Losses: logitcrossentropy
+```
+
+Loss functions for supervised learning typically expect as inputs a true target `y` and a prediction `ŷ`,
+typically passed as arrays of size `num_target_features x num_examples_in_batch`. 
+In Flux's convention, the order of the arguments is the following:
 
 ```julia
 loss(ŷ, y)
@@ -14,32 +22,16 @@ Most loss functions in Flux have an optional argument `agg`, denoting the type o
 batch:
 
 ```julia
-loss(ŷ, y)                         # defaults to `mean`
-loss(ŷ, y, agg=sum)                # use `sum` for reduction
-loss(ŷ, y, agg=x->sum(x, dims=2))  # partial reduction
-loss(ŷ, y, agg=x->mean(w .* x))    # weighted mean
-loss(ŷ, y, agg=identity)           # no aggregation.
+loss(ŷ, y)                             # defaults to `mean`
+loss(ŷ, y, agg = sum)                  # use `sum` for reduction
+loss(ŷ, y, agg = x -> sum(x, dims=2))  # partial reduction
+loss(ŷ, y, agg = x -> mean(w .* x))    # weighted mean
+loss(ŷ, y, agg = identity)             # no aggregation.
 ```
 
 ## Losses Reference
 
-```@docs
-Flux.Losses.mae
-Flux.Losses.mse
-Flux.Losses.msle
-Flux.Losses.huber_loss
-Flux.Losses.label_smoothing
-Flux.Losses.crossentropy
-Flux.Losses.logitcrossentropy
-Flux.Losses.binarycrossentropy
-Flux.Losses.logitbinarycrossentropy
-Flux.Losses.kldivergence
-Flux.Losses.poisson_loss
-Flux.Losses.hinge_loss
-Flux.Losses.squared_hinge_loss
-Flux.Losses.dice_coeff_loss
-Flux.Losses.tversky_loss
-Flux.Losses.binary_focal_loss
-Flux.Losses.focal_loss
-Flux.Losses.siamese_contrastive_loss
+```@autodocs
+Modules = [Flux.Losses]
+Pages   = ["functions.jl"]
 ```

src/Flux.jl

@@ -51,9 +51,8 @@ include("outputsize.jl")
 include("data/Data.jl")
 using .Data
 
-
 include("losses/Losses.jl")
-using .Losses # TODO: stop importing Losses in Flux's namespace in v0.12
+using .Losses # TODO: stop importing Losses in Flux's namespace in v0.14?
 
 include("deprecations.jl")