Montecarlo dropout uncertainty estimation#258
Conversation
lazarusA
commented
Apr 8, 2026
lazarusA
commented
- generate/save evaluation samples
- get statistics
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Monte Carlo (MC) Dropout evaluation functionality, adding the Functors dependency to support model layer traversal. The review identifies several critical issues: the current use of compute_loss fails to capture predictions when train_mode is enabled, samples are lost when no file path is provided for in-memory evaluation, and the implementation suffers from significant I/O overhead due to opening and closing the JLD2 file within the sampling loop.
| loss_k, _, ŷ_k = compute_loss( | ||
| ghm, ps, st_train, | ||
| (x, (y, y_no_nan)), | ||
| logging = LoggingLoss( | ||
| train_mode = true, | ||
| loss_types = loss_types, | ||
| training_loss = training_loss, | ||
| extra_loss = extra_loss, | ||
| agg = agg | ||
| ) | ||
| ) |
There was a problem hiding this comment.
The compute_loss function returns an empty NamedTuple for the stats field (the third return value) when logging.train_mode is set to true. Since this loop explicitly sets train_mode = true to keep dropout active, ŷ_k will be empty, and the model's predictions will not be captured. This prevents the MC Dropout from collecting the necessary samples for uncertainty estimation. You should consider calling the model directly to obtain predictions while in training mode, or adjusting compute_loss to return predictions even when train_mode is true.
| for k in 1:n_samples | ||
| loss_k, _, ŷ_k = compute_loss( | ||
| ghm, ps, st_train, | ||
| (x, (y, y_no_nan)), | ||
| logging = LoggingLoss( | ||
| train_mode = true, | ||
| loss_types = loss_types, | ||
| training_loss = training_loss, | ||
| extra_loss = extra_loss, | ||
| agg = agg | ||
| ) | ||
| ) | ||
| _store_sample(file_path, train_or_val_name, ŷ_k, loss_k, k) | ||
| end |
There was a problem hiding this comment.
When file_path is nothing, the samples generated in the loop are not collected or returned. The _store_sample function returns the prediction and loss, but these values are ignored by the loop, and evaluate_mc_dropout returns nothing. If the intention is to allow in-memory evaluation, you should accumulate the results in a list and return them at the end of the function.
|
|
||
|
|
||
| function _store_sample(file_path::String, name, ŷ, loss, sample) | ||
| return jldopen(file_path, "a+") do file |
There was a problem hiding this comment.
Opening and closing the JLD2 file with jldopen in append mode ("a+") inside the loop is inefficient. For a large number of samples, this will cause significant I/O overhead. It is recommended to open the file once before the loop starts and pass the open file handle to the storage function.
|
MC dropout does not make sense with globally estimated parameters, let's go for ensembles for now |
2 participants
