From e0c1c0e057dd9bf030f7289ad16283536f3313f4 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 17 Oct 2019 11:01:28 -0400 Subject: [PATCH] Fix problem in crossentropy breaking GPU compilation --- src/layers/stateless.jl | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/layers/stateless.jl b/src/layers/stateless.jl index 4c216672..6d710c6b 100644 --- a/src/layers/stateless.jl +++ b/src/layers/stateless.jl @@ -4,10 +4,20 @@ using NNlib: logsoftmax, logσ mse(ŷ, y) = sum((ŷ .- y).^2) * 1 // length(y) -function crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight = 1) - -sum(y .* log.(ŷ) .* weight) * 1 // size(y, 2) +function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Nothing) + return -sum(y .* log.(ŷ)) * 1 // size(y, 2) end +function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Number) + return -sum(y .* log.(ŷ)) .* weight * 1 // size(y, 2) +end + +function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::AbstractVector) + return -sum(y .* log.(ŷ) .* weight) * 1 // size(y, 2) +end + +crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight=nothing) = _crossentropy(ŷ, y, weight) + function logitcrossentropy(logŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight = 1) return -sum(y .* logsoftmax(logŷ) .* weight) * 1 // size(y, 2) end