Hi all,
I had a quick question – seemingly trivial but haven’t managed to figure it out after an hour of trying…
What’s the best way to log the training loss and ouput say to a simple text file for visualization using Flux. Consider for instance a template example using NeuralOperators.jl
slight modified, say
function get_data_don(; n=2048, Δsamples=2^3, grid_size=div(2^13, Δsamples))
file = matopen(joinpath(datadep"Burgers", "burgers_data_R10.mat"))
x_data = collect(read(file, "a")[1:n, 1:Δsamples:end])
y_data = collect(read(file, "u")[1:n, 1:Δsamples:end])
close(file)
return x_data, y_data
end
function train_don(; n=300, cuda=true, learning_rate=0.001, epochs=400)
if cuda && has_cuda()
@info "Training on GPU"
device = gpu
else
@info "Training on CPU"
device = cpu
end
x, y = get_data_don(n=n)
xtrain = x[1:280, :]'
ytrain = y[1:280, :]
xval = x[end-19:end, :]' |> device
yval = y[end-19:end, :] |> device
grid = collect(range(0, 1, length=1024)') |> device
opt = ADAM(learning_rate)
m = DeepONet((1024,1024,1024), (1,1024,1024), gelu, gelu) |> device
loss(X, y, sensor) = Flux.Losses.mse(m(X, sensor), y)
evalcb() = @show(loss(xval, yval, grid))
data = [(xtrain, ytrain, grid)] |> device
# Flux's default training loop
Flux.@epochs epochs Flux.train!(loss, params(m), data, opt, cb=evalcb)
# What we need instead is a custom training loop
global loss_vector = Vector{Float32}();
my_custom_train!(loss, ps, data, opt) # BUT THIS DOESN'T WORK
ỹ = m(xval |> device, grid |> device)
diffvec = vec(abs.(cpu(yval) .- cpu(ỹ)))
mean_diff = sum(diffvec)/length(diffvec)
return mean_diff
end
which is along the lines of what the docs on Flux indicate, namely,
function logging_callback(loss)
global loss_vector;
push!(loss_vector, loss);
end
function my_custom_train!(loss, ps, data, opt)
global loss_vector
local training_loss
ps = Flux.Params(ps) # Simply using Params throws an error
for d ∈ data
gs = gradient(ps) do
training_loss = loss(d...)
return training_loss
end
loggin_callback(training_loss)
update!(opt, ps, gs)
end
end
This is in essence similar to Flux: Custom Training + Logging - #4 by contradict but I couldn’t understand how I could simply append the loss to a vector and then write to a simple txt file for visualization downstream… Any pointers?
Thanks,