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))