I want to setup my model to use cross entropy loss in Flux.jl. How can I do this and where would I pass the loss function itself?
1 Answers
Flux.jl provides a build in module with many common loss function via the Flux.Losses module which you can access by doing using Flux.Losses. There is a built in cross entropy loss function in the module which can be used as follows:
julia> y_label = Flux.onehotbatch([0, 1, 2, 1, 0], 0:2)
3×5 Flux.OneHotArray{3,2,Vector{UInt32}}:
1 0 0 0 1
0 1 0 1 0
0 0 1 0 0
julia> y_model = softmax(reshape(-7:7, 3, 5) .* 1f0)
3×5 Matrix{Float32}:
0.0900306 0.0900306 0.0900306 0.0900306 0.0900306
0.244728 0.244728 0.244728 0.244728 0.244728
0.665241 0.665241 0.665241 0.665241 0.665241
julia> sum(y_model; dims=1)
1×5 Matrix{Float32}:
1.0 1.0 1.0 1.0 1.0
julia> Flux.crossentropy(y_model, y_label)
1.6076053f0
You can find the complete Flux.crossentropy function definition here: https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.crossentropy
After you define the loss function, you can pass it to the built in train function: Flux.train!(loss, params(model), data, opt) or use it in your custom training loop.
Related