FluxML · CarloLucibello · Dec 28, 2022 · Dec 28, 2022 · Dec 28, 2022 · Dec 28, 2022
diff --git a/src/losses/functions.jl b/src/losses/functions.jl
@@ -20,7 +20,7 @@
 """
 function mae(ŷ, y; agg = mean)
   _check_sizes(ŷ, y)
-  agg(abs.(ŷ .- y))
+  return agg(abs.(ŷ .- y))
 end
 
 """
@@ -44,7 +44,7 @@
 """
 function mse(ŷ, y; agg = mean)
   _check_sizes(ŷ, y)
-  agg(abs2.(ŷ .- y))
+  return agg(abs2.(ŷ .- y))
 end
 
 """
@@ -93,12 +93,11 @@
 ```
 """
 function huber_loss(ŷ, y; agg = mean, δ = ofeltype(ŷ, 1))
-   _check_sizes(ŷ, y)
-   abs_error = abs.(ŷ .- y)
-   #TODO: remove dropgrad when Zygote can handle this function with CuArrays
-   temp = Zygote.dropgrad(abs_error .<  δ)
-   x = ofeltype(ŷ, 0.5)
-   agg(((abs_error .^ 2) .* temp) .* x .+ δ * (abs_error .- x * δ) .* (1 .- temp))
+  _check_sizes(ŷ, y)
+  abs_error = abs.(ŷ .- y)
+  temp = abs_error .<  δ
+  x = ofeltype(ŷ, 0.5)
+  return agg(((abs_error .^ 2) .* temp) .* x .+ δ * (abs_error .- x * δ) .* (1 .- temp))
 end
 
 """
@@ -153,17 +152,17 @@
 ```
 """
 function label_smoothing(y::Union{AbstractArray,Number}, α::Number; dims::Int = 1)
-    if !(0 < α < 1)
-        throw(ArgumentError("α must be between 0 and 1"))
-    end
-    if dims == 0
-        y_smoothed = y .* (1 - α) .+ α*1//2
-    elseif dims == 1
-        y_smoothed = y .* (1 - α) .+ α* 1 // size(y, 1)
-    else
-        throw(ArgumentError("`dims` should be either 0 or 1"))
-    end
-    return y_smoothed
+  if !(0 < α < 1)
+    throw(ArgumentError("α must be between 0 and 1"))
+  end
+  if dims == 0
+    y_smoothed = y .* (1 - α) .+ α*1//2
+  elseif dims == 1
+    y_smoothed = y .* (1 - α) .+ α* 1 // size(y, 1)
+  else
+    throw(ArgumentError("`dims` should be either 0 or 1"))
+  end
+  return y_smoothed
 end
 
 """
@@ -224,7 +223,7 @@
 """
 function crossentropy(ŷ, y; dims = 1, agg = mean, ϵ = epseltype(ŷ))
   _check_sizes(ŷ, y)
-  agg(.-sum(xlogy.(y, ŷ .+ ϵ); dims = dims))
+  return agg(.-sum(xlogy.(y, ŷ .+ ϵ); dims))
 end
 
 """
@@ -263,7 +262,7 @@
 """
 function logitcrossentropy(ŷ, y; dims = 1, agg = mean)
   _check_sizes(ŷ, y)
-  agg(.-sum(y .* logsoftmax(ŷ; dims = dims); dims = dims))
+  return agg(.-sum(y .* logsoftmax(ŷ; dims); dims))
 end
 
 """
@@ -312,7 +311,7 @@
 """
 function binarycrossentropy(ŷ, y; agg = mean, ϵ = epseltype(ŷ))
   _check_sizes(ŷ, y)
-  agg(@.(-xlogy(y, ŷ + ϵ) - xlogy(1 - y, 1 - ŷ + ϵ)))
+  return agg(@.(-xlogy(y, ŷ + ϵ) - xlogy(1 - y, 1 - ŷ + ϵ)))
 end
 
 """
@@ -342,7 +341,7 @@
 """
 function logitbinarycrossentropy(ŷ, y; agg = mean)
   _check_sizes(ŷ, y)
-  agg(@.((1 - y) * ŷ - logσ(ŷ)))
+  return agg(@.((1 - y) * ŷ - logσ(ŷ)))
 end
 
 """
@@ -382,8 +381,8 @@
 """
 function kldivergence(ŷ, y; dims = 1, agg = mean, ϵ = epseltype(ŷ))
   _check_sizes(ŷ, y)
-  entropy = agg(sum(xlogx.(y), dims = dims))
-  cross_entropy = crossentropy(ŷ, y; dims = dims, agg = agg, ϵ = ϵ)
+  entropy = agg(sum(xlogx.(y); dims))
+  cross_entropy = crossentropy(ŷ, y; dims, agg, ϵ)
   return entropy + cross_entropy
 end
 
@@ -407,7 +406,7 @@
 """
 function poisson_loss(ŷ, y; agg = mean)
   _check_sizes(ŷ, y)
-  agg(ŷ .- xlogy.(y, ŷ))
+  return agg(ŷ .- xlogy.(y, ŷ))
 end
 
 """
@@ -442,7 +441,7 @@
 """
 function hinge_loss(ŷ, y; agg = mean)
   _check_sizes(ŷ, y)
-  agg(max.(0, 1 .- ŷ .* y))
+  return agg(max.(0, 1 .- ŷ .* y))
 end
 
 """
@@ -477,7 +476,7 @@
 """
 function squared_hinge_loss(ŷ, y; agg = mean)
   _check_sizes(ŷ, y)
-  agg((max.(0, 1 .- ŷ .* y)) .^ 2)
+  return agg((max.(0, 1 .- ŷ .* y)) .^ 2)
 end
 
 """
@@ -503,7 +502,7 @@
 """
 function dice_coeff_loss(ŷ, y; smooth = ofeltype(ŷ, 1.0))
   _check_sizes(ŷ, y)
-  1 - (2 * sum(y .* ŷ) + smooth) / (sum(y .^ 2) + sum(ŷ .^ 2) + smooth) #TODO agg
+  return 1 - (2 * sum(y .* ŷ) + smooth) / (sum(y .^ 2) + sum(ŷ .^ 2) + smooth) #TODO agg
 end
 
 """
@@ -518,11 +517,11 @@
 
 """
 function tversky_loss(ŷ, y; β = ofeltype(ŷ, 0.7))
-    _check_sizes(ŷ, y)
-    #TODO add agg
-    num = sum(y .* ŷ) + 1
-    den = sum(y .* ŷ + β * (1 .- y) .* ŷ + (1 - β) * y .* (1 .- ŷ)) + 1
-    1 - num / den
+  _check_sizes(ŷ, y)
+  #TODO add agg
+  num = sum(y .* ŷ) + 1
+  den = sum(y .* ŷ + β * (1 .- y) .* ŷ + (1 - β) * y .* (1 .- ŷ)) + 1
+  return 1 - num / den
 end
 
 """
@@ -554,13 +553,13 @@
 ```
 """
 function binary_focal_loss(ŷ, y; agg=mean, γ=2, ϵ=epseltype(ŷ))
-    _check_sizes(ŷ, y)
-    ŷ = ŷ .+ ϵ
-    p_t = y .* ŷ  + (1 .- y) .* (1 .- ŷ)
-    ce = -log.(p_t)
-    weight = (1 .- p_t) .^ γ
-    loss = weight .* ce
-    agg(loss)
+  _check_sizes(ŷ, y)
+  ŷ = ŷ .+ ϵ
+  p_t = y .* ŷ  + (1 .- y) .* (1 .- ŷ)
+  ce = -log.(p_t)
+  weight = (1 .- p_t) .^ γ
+  loss = weight .* ce
+  return agg(loss)
 end
 
 """
@@ -597,10 +596,10 @@
 See also: [`Losses.binary_focal_loss`](@ref) for binary (not one-hot) labels
 
 """
-function focal_loss(ŷ, y; dims=1, agg=mean, γ=2, ϵ=epseltype(ŷ))
-    _check_sizes(ŷ, y)
-    ŷ = ŷ .+ ϵ
-    agg(sum(@. -y * (1 - ŷ)^γ * log(ŷ); dims=dims))
+function focal_loss(ŷ, y; dims = 1, agg = mean, γ = 2, ϵ = epseltype(ŷ))
+  _check_sizes(ŷ, y)
+  ŷ = ŷ .+ ϵ
+  return agg(sum(@. -y * (1 - ŷ)^γ * log(ŷ); dims))
 end
 
 """
@@ -625,9 +624,9 @@
 ```
 """
 function siamese_contrastive_loss(ŷ, y; agg = mean, margin::Real = 1)
-    _check_sizes(ŷ, y)
-    margin < 0 && throw(DomainError(margin, "Margin must be non-negative"))
-    return agg(@. (1 - y) * ŷ^2 + y * max(0, margin - ŷ)^2)
+  _check_sizes(ŷ, y)
+  margin < 0 && throw(DomainError(margin, "Margin must be non-negative"))
+  return agg(@. (1 - y) * ŷ^2 + y * max(0, margin - ŷ)^2)
 end
 
 ```@meta

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,5 +1,5 @@
 using Flux
-using Flux.Data
+using Flux: DataLoader
 using Flux: OneHotArray, OneHotMatrix, OneHotVector
 using Flux: params
 using Test

diff --git a/test/train.jl b/test/train.jl
@@ -24,6 +24,7 @@ using Random
 
   # Test direct use of Optimisers.jl rule, only really OK for `Descent`:
   @testset "without setup, $opt" for opt in [Descent(0.1), Optimisers.Descent(0.1), Optimisers.Adam()]
+    opt isa Optimisers.Adam && @info "should see a warning about discarded state"
     loss(m, x) = Flux.Losses.mse(w*x, m.weight*x .+ m.bias)
     model = (weight=copy(w2), bias=zeros(10), ignore=nothing)
     @test loss(model, rand(10, 10)) > 1

diff --git a/test/utils.jl b/test/utils.jl
@@ -385,18 +385,6 @@ end
     @test_skip typeof(l1.bias) === typeof(l2.bias)
   end
 
-  @testset "loadparams!" begin
-    pars(w, b) = [w, b]
-    pars(l) = pars(l.weight, l.bias)
-    pararray(m) = mapreduce(pars, vcat, m)
-    weights(m) = mapreduce(l -> [l.weight], vcat, m)
-    @testset "Bias type $bt" for bt in (Flux.zeros32, nobias)
-      m = dm(bt)
-      Flux.loadparams!(m, params(m))
-      testdense(m, bt)
-    end
-  end
-
   @testset "loadmodel!(dst, src)" begin
     m1 = Chain(Dense(10, 5), Dense(5, 2, relu))
     m2 = Chain(Dense(10, 5), Dense(5, 2))