make loss(f,x,y) work

Author: mcabbott

docs/src/models/losses.md

@@ -4,21 +4,27 @@ Flux provides a large number of common loss functions used for training machine
 They are grouped together in the `Flux.Losses` module.
 
 Loss functions for supervised learning typically expect as inputs a target `y`, and a prediction `ŷ` from your model.
-In Flux's convention, the order of the arguments is the following
+In Flux's convention, the target is the last argumemt:
 
 ```julia
 loss(ŷ, y)
 ```
 
+All loss functions have a method which takes the model as the first argument, and calculates the prediction `ŷ = model(x)`.
+This is convenient for [`train!`](@ref Flux.train)`(loss, model, [(x,y), (x2,y2), ...], opt)`:
+
+```julia
+loss(model, x, y) = loss(model(x), y)
+```
+
 Most loss functions in Flux have an optional argument `agg`, denoting the type of aggregation performed over the
 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=sum)                # use `sum` instead
 loss(ŷ, y, agg=x->mean(w .* x))    # weighted mean
-loss(ŷ, y, agg=identity)           # no aggregation.
+loss(ŷ, y, agg=x->sum(x, dims=2))  # partial reduction, returns an array
 ```
 
 ### Function listing

src/losses/Losses.jl

@@ -1,3 +1,28 @@
+"""
+    Flux.Losses
+    
+This sub-module contains many loss functions, all of which accept two arguments,
+with the model output as the fist argument: `loss(model(x), y)`.
+It also contains a few related utilities, such as `label_smoothing`.
+The complete list of exports is:
+
+    label_smoothing,
+    mse, mae, msle,
+    crossentropy,
+    logitcrossentropy,
+    binarycrossentropy,
+    logitbinarycrossentropy,
+    kldivergence,
+    huber_loss,
+    tversky_loss,
+    dice_coeff_loss,
+    poisson_loss,
+    hinge_loss,
+    squared_hinge_loss,
+    binary_focal_loss,
+    focal_loss,
+    siamese_contrastive_loss
+"""
 module Losses
 
 using Statistics
@@ -9,8 +34,8 @@ using CUDA
 using NNlib: logsoftmax, logσ, ctc_loss, ctc_alpha, ∇ctc_loss
 import Base.Broadcast: broadcasted
 
-export mse, mae, msle,
-    label_smoothing,
+export label_smoothing,
+    mse, mae, msle,
     crossentropy, logitcrossentropy,
     binarycrossentropy, logitbinarycrossentropy,
     kldivergence,
@@ -19,9 +44,33 @@ export mse, mae, msle,
     dice_coeff_loss,
     poisson_loss,
     hinge_loss, squared_hinge_loss,
-    binary_focal_loss, focal_loss, siamese_contrastive_loss
+    binary_focal_loss, focal_loss,
+    siamese_contrastive_loss
 
 include("utils.jl")
 include("functions.jl")
 
+for loss in Symbol.([
+  mse, mae, msle,
+  crossentropy, logitcrossentropy,
+  binarycrossentropy, logitbinarycrossentropy,
+  kldivergence,
+  huber_loss,
+  tversky_loss,
+  dice_coeff_loss,
+  poisson_loss,
+  hinge_loss, squared_hinge_loss,
+  binary_focal_loss, focal_loss,
+  siamese_contrastive_loss,
+  ])
+  @eval begin
+    """
+        $($loss)(model, x, y)
+  
+    This method calculates `ŷ = model(x)`. Accepts the same keyword arguments.
+    """
+    $loss(f, x::AbstractArray, y::AbstractArray; kw...) = $loss(f(x), y; kw...)
+  end
+end
+
 end #module

test/losses.jl

@@ -248,3 +248,12 @@ end
   @test_throws DomainError(-0.5, "Margin must be non-negative") Flux.siamese_contrastive_loss(ŷ1, y1, margin = -0.5)
   @test_throws DomainError(-1, "Margin must be non-negative") Flux.siamese_contrastive_loss(ŷ, y, margin = -1)
 end
+
+@testset "3-arg methods" begin
+  @testset for loss in ALL_LOSSES
+    fun(x) = x[1:2]
+    x = rand(3)
+    y = rand(2)
+    @test loss(fun, x, y) == loss(fun(x), y)
+  end
+end