How-to Finetune

Before we start

This tutorial is rendered from a Jupyter notebook that is hosted on GitHub. If you want to run the code yourself, you can find the notebook and configuration files here.
To be able to run this notebook locally, you need to download the publicly available CAMELS US rainfall-runoff dataset and a publicly available extensions for hourly forcing and streamflow data. See the Data Prerequisites Tutorial for a detailed description on where to download the data and how to structure your local dataset folder. Note the special section with additional requirements for this tutorial.

This tutorial shows how to adapt a pretrained model to a different, eventually much smaller dataset, a concept called finetuning. Finetuning is well-established in machine learning and thus nothing new. Generally speaking, the idea is to use a (very) large and diverse dataset to learn a general understanding of the underlying problem first and then, in a second step, adapt this general model to the target data. Usually, especially if the available target data is limited, pretraining plus finetuning yields (much) better results than only considering the final target data.

The connection to hydrology is the following: Often, researchers or operators are only interested in a single basin. However, considering that a Deep Learning (DL) model has to learn all (physical) process understanding from the available training data, it might be understandable that the data records of a single basin might not be enough (see e.g. the presentation linked at this EGU’20 abstract)

This is were we apply the concept of pretraining and finetuning: First, we train a DL model (e.g. an LSTM) with a large and diverse, multi-basin dataset (e.g. CAMELS) and then finetune this model to our basin of interest. Everything you need is available in the NeuralHydrology package and in this notebook we will give you an overview of how to actually do it.

Note: Finetuning can be a tedious task and is usually very sensitive to the learning rate as well as the number of epochs used for finetuning. One reason is that the pretrained models are usually quite large. In fact, most often they are much larger than what would be possible to train for just a single basin. So during finetuning, we have to make sure that this large capacity is not negatively impacting our model results. Common approaches are to a) only allow parts of the model to be adapted during finetuning and/or b) to train with a much lower learning rate. So far, no publication was published that presents a universally working approach for finetuning in hydrology. So be aware that the results may vary and you might need to invest some time before finding a good strategy. However, in our experience it was always possible to get better results with finetuning than without.

To summarize: If you are interested in getting the best-performing Deep Learning model for a single basin, pretraining on a large and diverse dataset, followed by finetuning the pretrained model on your target basin is the way to go.

[1]:

# Imports
from pathlib import Path

import pandas as pd
import torch
from neuralhydrology.nh_run import start_run, eval_run, finetune

Pretraining

In the first step, we need to pretrain our model on a large and possibly diverse dataset. Our target basin does not necessarily have to be a part of this dataset, but usually it should be better to include it.

For the sake of the demonstration, we will train an LSTM on the CAMELS US dataset and then finetune this model to a random basin. Note that it is possible to use other inputs during pretraining and finetuning, if additional embedding layers (before the LSTM) are used, which we will ignore for now. Furthermore, we will concentrate only on demonstrating the “how-to” rather than striving for best-possible performance. To save time and energy, we will only pretrain the model for a small number of epochs. When striving for the best possible performance, you should make sure that you pretrain the model as best as possible, before starting to finetune.

We will stick closely to the model and experimental setup from Kratzert et al. (2019). To summarize: - A single LSTM layer with a hidden size of 128. - Input sequences are 365 days and the prediction is made at the last timestep. - We will use the same CAMELS attributes, as in the publication mentioned above, as additional inputs at every time step so that the model can learn different hydrological behaviors depending on the catchment properties.

For more details, take a look at the config print-out below.

Note - The config file assumes that the CAMELS US dataset is stored under data/CAMELS_US (relative to the main directory of this repository) or a symbolic link exists at this location. Make sure that this folder contains the required subdirectories basin_mean_forcing, usgs_streamflow and camels_attributes_v2.0. If your data is stored at a different location and you can’t or don’t want to create a symbolic link, you will need to change the data_dir argument in the 531_basins.yml config file that is located in the same directory as this notebook. - By default, the config (531_basins.yml) assumes that you have a CUDA-capable NVIDIA GPU (see config argument device). In case you don’t have any or you have one but want to train on the CPU, you can either change the config argument to device: cpu or pass gpu=-1 to the start_run() function. Please note that training such a model on such a large dataset on CPU takes a very long time.

[2]:

# by default we assume that you have at least one CUDA-capable NVIDIA GPU
if torch.cuda.is_available():
    start_run(config_file=Path("531_basins.yml"))

# fall back to CPU-only mode
else:
    start_run(config_file=Path("531_basins.yml"), gpu=-1)

2022-01-09 13:54:00,506: Logging to /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/output.log initialized.
2022-01-09 13:54:00,507: ### Folder structure created at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400
2022-01-09 13:54:00,507: ### Run configurations for cudalstm_531_basins
2022-01-09 13:54:00,507: experiment_name: cudalstm_531_basins
2022-01-09 13:54:00,508: run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400
2022-01-09 13:54:00,508: train_basin_file: 531_basin_list.txt
2022-01-09 13:54:00,509: validation_basin_file: 531_basin_list.txt
2022-01-09 13:54:00,509: test_basin_file: 531_basin_list.txt
2022-01-09 13:54:00,509: train_start_date: 1999-10-01 00:00:00
2022-01-09 13:54:00,510: train_end_date: 2008-09-30 00:00:00
2022-01-09 13:54:00,510: validation_start_date: 1980-10-01 00:00:00
2022-01-09 13:54:00,510: validation_end_date: 1989-09-30 00:00:00
2022-01-09 13:54:00,510: test_start_date: 1989-10-01 00:00:00
2022-01-09 13:54:00,511: test_end_date: 1999-09-30 00:00:00
2022-01-09 13:54:00,511: seed: 123
2022-01-09 13:54:00,511: device: cuda:0
2022-01-09 13:54:00,511: validate_every: 1
2022-01-09 13:54:00,511: validate_n_random_basins: 531
2022-01-09 13:54:00,512: metrics: ['NSE']
2022-01-09 13:54:00,512: model: cudalstm
2022-01-09 13:54:00,513: head: regression
2022-01-09 13:54:00,513: hidden_size: 128
2022-01-09 13:54:00,513: initial_forget_bias: 3
2022-01-09 13:54:00,514: output_dropout: 0.4
2022-01-09 13:54:00,514: output_activation: linear
2022-01-09 13:54:00,514: optimizer: Adam
2022-01-09 13:54:00,515: loss: NSE
2022-01-09 13:54:00,515: learning_rate: {0: 0.001, 1: 0.0005}
2022-01-09 13:54:00,515: batch_size: 256
2022-01-09 13:54:00,516: epochs: 3
2022-01-09 13:54:00,516: clip_gradient_norm: 1
2022-01-09 13:54:00,516: predict_last_n: 1
2022-01-09 13:54:00,517: seq_length: 365
2022-01-09 13:54:00,517: num_workers: 8
2022-01-09 13:54:00,518: log_interval: 5
2022-01-09 13:54:00,518: log_tensorboard: True
2022-01-09 13:54:00,518: save_weights_every: 1
2022-01-09 13:54:00,519: save_validation_results: True
2022-01-09 13:54:00,519: dataset: camels_us
2022-01-09 13:54:00,519: data_dir: ../../data/CAMELS_US
2022-01-09 13:54:00,520: forcings: ['daymet']
2022-01-09 13:54:00,520: dynamic_inputs: ['prcp(mm/day)', 'srad(W/m2)', 'tmax(C)', 'tmin(C)', 'vp(Pa)']
2022-01-09 13:54:00,521: target_variables: ['QObs(mm/d)']
2022-01-09 13:54:00,521: static_attributes: ['elev_mean', 'slope_mean', 'area_gages2', 'frac_forest', 'lai_max', 'lai_diff', 'gvf_max', 'gvf_diff', 'soil_depth_pelletier', 'soil_depth_statsgo', 'soil_porosity', 'soil_conductivity', 'max_water_content', 'sand_frac', 'silt_frac', 'clay_frac', 'carbonate_rocks_frac', 'geol_permeability', 'p_mean', 'pet_mean', 'aridity', 'frac_snow', 'high_prec_freq', 'high_prec_dur', 'low_prec_freq', 'low_prec_dur']
2022-01-09 13:54:00,521: number_of_basins: 531
2022-01-09 13:54:00,522: train_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/train_data
2022-01-09 13:54:00,522: img_log_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/img_log
2022-01-09 13:54:00,571: ### Device cuda:0 will be used for training
2022-01-09 13:54:02,440: Loading basin data into xarray data set.
100%|██████████| 531/531 [00:30<00:00, 17.37it/s]
2022-01-09 13:54:33,325: Calculating target variable stds per basin
100%|██████████| 531/531 [00:00<00:00, 1808.07it/s]
2022-01-09 13:54:33,678: Create lookup table and convert to pytorch tensor
100%|██████████| 531/531 [00:10<00:00, 51.85it/s]
2022-01-09 13:54:44,231: Setting learning rate to 0.0005
# Epoch 1: 100%|██████████| 6821/6821 [03:15<00:00, 34.95it/s, Loss: 0.0021]
2022-01-09 13:57:59,411: Epoch 1 average loss: 0.036746779475167225
# Validation:  29%|██▉       | 156/531 [00:41<02:23,  2.62it/s]2022-01-09 13:58:41,119: The following basins had not enough valid target values to calculate a standard deviation: 02427250. NSE loss values for this basin will be NaN.
# Validation:  79%|███████▊  | 418/531 [01:50<00:30,  3.71it/s]2022-01-09 13:59:50,108: The following basins had not enough valid target values to calculate a standard deviation: 09484600. NSE loss values for this basin will be NaN.
# Validation: 100%|██████████| 531/531 [02:20<00:00,  3.78it/s]
2022-01-09 14:00:19,838: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/validation/model_epoch001/validation_results.p
2022-01-09 14:00:19,846: Epoch 1 average validation loss: 0.02972 -- Median validation metrics: NSE: 0.62492
# Epoch 2: 100%|██████████| 6821/6821 [03:19<00:00, 34.22it/s, Loss: 0.0069]
2022-01-09 14:03:39,148: Epoch 2 average loss: 0.02329262817169138
# Validation: 100%|██████████| 531/531 [01:49<00:00,  4.87it/s]
2022-01-09 14:05:28,277: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/validation/model_epoch002/validation_results.p
2022-01-09 14:05:28,283: Epoch 2 average validation loss: 0.03369 -- Median validation metrics: NSE: 0.67444
# Epoch 3: 100%|██████████| 6821/6821 [03:18<00:00, 34.34it/s, Loss: 0.0106]
2022-01-09 14:08:46,896: Epoch 3 average loss: 0.021140944835559363
# Validation: 100%|██████████| 531/531 [01:46<00:00,  4.96it/s]
2022-01-09 14:10:33,952: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/validation/model_epoch003/validation_results.p
2022-01-09 14:10:33,957: Epoch 3 average validation loss: 0.02744 -- Median validation metrics: NSE: 0.68757

We end with an okay’ish model that should be enough for the purpose of this demonstration. Remember we only train for a limited number of epochs here.

Next, we’ll load the validation results into memory so we can select a basin to demonstrate how to finetune based on the model performance. Since the folder name is created dynamically (including the date and time of the start of the run) you will need to change the run_dir argument according to your local directory name.

Here, we will select a random basin from the lower 50% of the NSE distribution, i.e. a basin where the NSE is below the median NSE. Usually, you’ll see better performance gains for basins with lower model performance than for those where the base model is already really good.

[16]:

# Load validation results from the last epoch
run_dir = Path("runs/cudalstm_531_basins_0901_135400/")
df = pd.read_csv(run_dir / "validation" / "model_epoch003" / "validation_metrics.csv", dtype={'basin': str})
df = df.set_index('basin')

# Compute the median NSE from all basins, where discharge observations are available for that period
print(f"Median NSE of the validation period {df['NSE'].median():.3f}")

# Select a random basins from the lower 50% of the NSE distribution
basin = df.loc[df["NSE"] < df["NSE"].median()].sample(n=1).index[0]

print(f"Selected basin: {basin} with an NSE of {df.loc[df.index == basin, 'NSE'].values[0]:.3f}")

Median NSE of the validation period 0.688
Selected basin: 02112360 with an NSE of 0.501

Finetuning

Next, we will show how to perform finetuning for the basin selected above, based on the model we just trained. The function to use is finetune from neuralhydrology.nh_run if you want to train from within a script or notebook. If you want to start finetuning from the command line, you can also call the nh-run utility with the finetune argument, instead of e.g. train or evaluate.

The only thing required, similar to the model training itself, is a config file. This config however has slightly different requirements to a normal model config and works slightly different: - The config has to contain the following two arguments: - base_run_dir: The path to the directory of the pre-trained model. - finetune_modules: Which parts of the pre-trained model you want to finetune. Check the documentation of each model class for a list of all possible parts. Often only parts, e.g. the output layer, are trained during finetuning and the rest is kept fixed. There is no general rule of thumb and most likely you will have to try both. - Any additional argument contained in this config will overwrite the config argument of the pre-trained model. Everything not specified will be taken from the pre-trained model. That is, you can e.g. specify a new basin file in the finetuning config (by train_basin_file) to finetune the pre-trained model on a different set of basins, or even just a single basin as we will do in this notebook. You can also change the learning rate, loss function, evaluation metrics and so on. The only thing you can not change are arguments that change the model architecture (e.g. model, hidden_size etc.), because this leads to errors when you try to load the pre-trained weights into the initialized model.

Let’s have a look at the finetune.yml config that we prepared for this tutorial (you can find the file in the same directory as this notebook).

[17]:

!cat finetune.yml

# --- Experiment configurations --------------------------------------------------------------------

# experiment name, used as folder name
experiment_name: cudalstm_531_basins_finetuned

# files to specify training, validation and test basins (relative to code root or absolute path)
train_basin_file: finetune_basin.txt
validation_basin_file: finetune_basin.txt
test_basin_file: finetune_basin.txt

# --- Training configuration -----------------------------------------------------------------------

# specify learning rates to use starting at specific epochs (0 is the initial learning rate)
learning_rate:
    0: 5e-4
    2: 5e-5

# Number of training epochs
epochs: 10

finetune_modules:
- head
- lstm
base_run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400

So out of the two arguments that are required, base_run_dir is still missing. We will add the argument from here and point at the directory of the model we just trained. Furthermore, we point to a new file for training, validation and testing, called finetune_basin.txt, which does not yet exist. We will create this file and add the basin we selected above as the only basin we want to use here. The rest are some changes to the learning rate and the number of training epochs as well as a new name. Also note that here, we train the full model, by selecting all model parts available for the CudaLSTM under finetune_modules.

[18]:

# Add the path to the pre-trained model to the finetune config
with open("finetune.yml", "a") as fp:
    fp.write(f"\nbase_run_dir: {run_dir.absolute()}")

# Create a basin file with the basin we selected above
with open("finetune_basin.txt", "w") as fp:
    fp.write(basin)

With that, we are ready to start the finetuning. As mentioned above, we have two options to start finetuning: 1. Call the finetune() function from a different Python script or a Jupyter Notebook with the path to the config. 2. Start the finetuning from the command line by calling

nh-run finetune --config-file /path/to/config.yml

Here, we will use the first option.

[19]:

finetune(Path("finetune.yml"))

2022-01-09 14:15:48,842: Logging to /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/output.log initialized.
2022-01-09 14:15:48,843: ### Folder structure created at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548
2022-01-09 14:15:48,843: ### Start finetuning with pretrained model stored in /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400
2022-01-09 14:15:48,844: ### Run configurations for cudalstm_531_basins_finetuned
2022-01-09 14:15:48,844: batch_size: 256
2022-01-09 14:15:48,844: clip_gradient_norm: 1
2022-01-09 14:15:48,844: commit_hash: b'37482bf\n'
2022-01-09 14:15:48,845: data_dir: ../../data/CAMELS_US
2022-01-09 14:15:48,845: dataset: camels_us
2022-01-09 14:15:48,845: device: cuda:0
2022-01-09 14:15:48,845: dynamic_inputs: ['prcp(mm/day)', 'srad(W/m2)', 'tmax(C)', 'tmin(C)', 'vp(Pa)']
2022-01-09 14:15:48,846: epochs: 10
2022-01-09 14:15:48,846: experiment_name: cudalstm_531_basins_finetuned
2022-01-09 14:15:48,846: forcings: ['daymet']
2022-01-09 14:15:48,847: head: regression
2022-01-09 14:15:48,847: hidden_size: 128
2022-01-09 14:15:48,847: img_log_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/img_log
2022-01-09 14:15:48,847: initial_forget_bias: 3
2022-01-09 14:15:48,848: learning_rate: {0: 0.0005, 2: 5e-05}
2022-01-09 14:15:48,848: log_interval: 5
2022-01-09 14:15:48,848: log_tensorboard: True
2022-01-09 14:15:48,849: loss: NSE
2022-01-09 14:15:48,849: metrics: ['NSE']
2022-01-09 14:15:48,849: model: cudalstm
2022-01-09 14:15:48,849: num_workers: 8
2022-01-09 14:15:48,849: number_of_basins: 1
2022-01-09 14:15:48,850: optimizer: Adam
2022-01-09 14:15:48,850: output_activation: linear
2022-01-09 14:15:48,850: output_dropout: 0.4
2022-01-09 14:15:48,850: package_version: 1.1.0
2022-01-09 14:15:48,851: predict_last_n: 1
2022-01-09 14:15:48,851: run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548
2022-01-09 14:15:48,851: save_validation_results: True
2022-01-09 14:15:48,851: save_weights_every: 1
2022-01-09 14:15:48,851: seed: 123
2022-01-09 14:15:48,852: seq_length: 365
2022-01-09 14:15:48,852: static_attributes: ['elev_mean', 'slope_mean', 'area_gages2', 'frac_forest', 'lai_max', 'lai_diff', 'gvf_max', 'gvf_diff', 'soil_depth_pelletier', 'soil_depth_statsgo', 'soil_porosity', 'soil_conductivity', 'max_water_content', 'sand_frac', 'silt_frac', 'clay_frac', 'carbonate_rocks_frac', 'geol_permeability', 'p_mean', 'pet_mean', 'aridity', 'frac_snow', 'high_prec_freq', 'high_prec_dur', 'low_prec_freq', 'low_prec_dur']
2022-01-09 14:15:48,853: target_variables: ['QObs(mm/d)']
2022-01-09 14:15:48,853: test_basin_file: finetune_basin.txt
2022-01-09 14:15:48,853: test_end_date: 1999-09-30 00:00:00
2022-01-09 14:15:48,853: test_start_date: 1989-10-01 00:00:00
2022-01-09 14:15:48,853: train_basin_file: finetune_basin.txt
2022-01-09 14:15:48,854: train_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/train_data
2022-01-09 14:15:48,854: train_end_date: 2008-09-30 00:00:00
2022-01-09 14:15:48,854: train_start_date: 1999-10-01 00:00:00
2022-01-09 14:15:48,854: validate_every: 1
2022-01-09 14:15:48,855: validate_n_random_basins: 531
2022-01-09 14:15:48,855: validation_basin_file: finetune_basin.txt
2022-01-09 14:15:48,855: validation_end_date: 1989-09-30 00:00:00
2022-01-09 14:15:48,855: validation_start_date: 1980-10-01 00:00:00
2022-01-09 14:15:48,856: finetune_modules: ['head', 'lstm']
2022-01-09 14:15:48,856: base_run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400
2022-01-09 14:15:48,856: is_finetuning: True
2022-01-09 14:15:48,856: is_continue_training: False
2022-01-09 14:15:48,857: ### Device cuda:0 will be used for training
2022-01-09 14:15:48,867: Starting training from checkpoint /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_0901_135400/model_epoch003.pt
2022-01-09 14:15:48,895: Loading basin data into xarray data set.
100%|██████████| 1/1 [00:00<00:00, 20.39it/s]
2022-01-09 14:15:48,947: Calculating target variable stds per basin
100%|██████████| 1/1 [00:00<00:00, 975.19it/s]
2022-01-09 14:15:48,951: Create lookup table and convert to pytorch tensor
100%|██████████| 1/1 [00:00<00:00, 68.31it/s]
# Epoch 1: 100%|██████████| 13/13 [00:00<00:00, 21.65it/s, Loss: 0.0156]
2022-01-09 14:15:49,685: Epoch 1 average loss: 0.019881912077275608
# Validation: 100%|██████████| 1/1 [00:00<00:00,  4.09it/s]
2022-01-09 14:15:49,936: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch001/validation_results.p
2022-01-09 14:15:49,937: Epoch 1 average validation loss: 0.03251 -- Median validation metrics: NSE: 0.58215
2022-01-09 14:15:49,937: Setting learning rate to 5e-05
# Epoch 2: 100%|██████████| 13/13 [00:00<00:00, 21.60it/s, Loss: 0.0074]
2022-01-09 14:15:50,541: Epoch 2 average loss: 0.01594447487821946
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.01it/s]
2022-01-09 14:15:50,715: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch002/validation_results.p
2022-01-09 14:15:50,716: Epoch 2 average validation loss: 0.03277 -- Median validation metrics: NSE: 0.57867
# Epoch 3: 100%|██████████| 13/13 [00:00<00:00, 22.47it/s, Loss: 0.0237]
2022-01-09 14:15:51,297: Epoch 3 average loss: 0.01625036971213726
# Validation: 100%|██████████| 1/1 [00:00<00:00,  5.98it/s]
2022-01-09 14:15:51,470: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch003/validation_results.p
2022-01-09 14:15:51,471: Epoch 3 average validation loss: 0.03320 -- Median validation metrics: NSE: 0.57281
# Epoch 4: 100%|██████████| 13/13 [00:00<00:00, 23.16it/s, Loss: 0.0071]
2022-01-09 14:15:52,034: Epoch 4 average loss: 0.015257955636256017
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.33it/s]
2022-01-09 14:15:52,198: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch004/validation_results.p
2022-01-09 14:15:52,199: Epoch 4 average validation loss: 0.03328 -- Median validation metrics: NSE: 0.57175
# Epoch 5: 100%|██████████| 13/13 [00:00<00:00, 23.93it/s, Loss: 0.0155]
2022-01-09 14:15:52,745: Epoch 5 average loss: 0.015026288751799326
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.11it/s]
2022-01-09 14:15:52,916: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch005/validation_results.p
2022-01-09 14:15:52,916: Epoch 5 average validation loss: 0.03254 -- Median validation metrics: NSE: 0.58182
# Epoch 6: 100%|██████████| 13/13 [00:00<00:00, 23.57it/s, Loss: 0.0706]
2022-01-09 14:15:53,470: Epoch 6 average loss: 0.015907029263102092
# Validation: 100%|██████████| 1/1 [00:00<00:00,  5.67it/s]
2022-01-09 14:15:53,652: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch006/validation_results.p
2022-01-09 14:15:53,653: Epoch 6 average validation loss: 0.03255 -- Median validation metrics: NSE: 0.58167
# Epoch 7: 100%|██████████| 13/13 [00:00<00:00, 23.42it/s, Loss: 0.0094]
2022-01-09 14:15:54,210: Epoch 7 average loss: 0.013952206784429459
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.17it/s]
2022-01-09 14:15:54,378: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch007/validation_results.p
2022-01-09 14:15:54,378: Epoch 7 average validation loss: 0.03184 -- Median validation metrics: NSE: 0.59141
# Epoch 8: 100%|██████████| 13/13 [00:00<00:00, 24.04it/s, Loss: 0.0092]
2022-01-09 14:15:54,921: Epoch 8 average loss: 0.015188249741465999
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.24it/s]
2022-01-09 14:15:55,086: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch008/validation_results.p
2022-01-09 14:15:55,087: Epoch 8 average validation loss: 0.03203 -- Median validation metrics: NSE: 0.58869
# Epoch 9: 100%|██████████| 13/13 [00:00<00:00, 22.61it/s, Loss: 0.0071]
2022-01-09 14:15:55,664: Epoch 9 average loss: 0.01575774785417777
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.35it/s]
2022-01-09 14:15:55,827: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch009/validation_results.p
2022-01-09 14:15:55,827: Epoch 9 average validation loss: 0.03222 -- Median validation metrics: NSE: 0.58619
# Epoch 10: 100%|██████████| 13/13 [00:00<00:00, 23.07it/s, Loss: 0.0170]
2022-01-09 14:15:56,393: Epoch 10 average loss: 0.015302825826578416
# Validation: 100%|██████████| 1/1 [00:00<00:00,  6.10it/s]
2022-01-09 14:15:56,562: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_531_basins_finetuned_0901_141548/validation/model_epoch010/validation_results.p
2022-01-09 14:15:56,563: Epoch 10 average validation loss: 0.03209 -- Median validation metrics: NSE: 0.58788

Looking at the validation result, we can see an increase of roughly 0.05 NSE.

Last but not least, we will compare the pre-trained and the finetuned model on the test period. For this, we will make use of the eval_run function from neuralhydrolgy.nh_run. Alternatively, you could evaluate both runs from the command line by calling

nh-run evaluate --run-dir /path/to/run_directory/

[20]:

eval_run(run_dir, period="test")

2022-01-09 14:16:09,586: Using the model weights from runs/cudalstm_531_basins_0901_135400/model_epoch003.pt
# Evaluation: 100%|██████████| 531/531 [02:09<00:00,  4.11it/s]
2022-01-09 14:18:18,959: Stored results at runs/cudalstm_531_basins_0901_135400/test/model_epoch003/test_results.p

Now we can call the eval_run() function as above, but pointing to the directory of the finetuned run. By default, this function evaluates the last checkpoint, which can be changed with the epoch argument. Here however, we use the default. Again, if you want to run this notebook locally, make sure to adapt the folder name of the finetune run.

[23]:

finetune_dir = Path("runs/cudalstm_531_basins_finetuned_0901_141548")
eval_run(finetune_dir, period="test")

2022-01-09 14:19:06,488: Using the model weights from runs/cudalstm_531_basins_finetuned_0901_141548/model_epoch010.pt
# Evaluation: 100%|██████████| 1/1 [00:00<00:00,  4.27it/s]
2022-01-09 14:19:06,726: Stored results at runs/cudalstm_531_basins_finetuned_0901_141548/test/model_epoch010/test_results.p

Now let’s look at the test period results of the pre-trained base model and the finetuned model for the basin that we chose above.

[24]:

# load test results of the base run
df_pretrained = pd.read_csv(run_dir / "test/model_epoch003/test_metrics.csv", dtype={'basin': str})
df_pretrained = df_pretrained.set_index("basin")

# load test results of the finetuned model
df_finetuned = pd.read_csv(finetune_dir / "test/model_epoch010/test_metrics.csv", dtype={'basin': str})
df_finetuned = df_finetuned.set_index("basin")

# extract basin performance
base_model_nse = df_pretrained.loc[df_pretrained.index == basin, "NSE"].values[0]
finetune_nse = df_finetuned.loc[df_finetuned.index == basin, "NSE"].values[0]
print(f"Basin {basin} base model performance: {base_model_nse:.3f}")
print(f"Performance after finetuning: {finetune_nse:.3f}")

Basin 02112360 base model performance: 0.303
Performance after finetuning: 0.580

So we see roughly the same performance increase in the test period (slightly higher), which is great. However, note that a) our base model was not optimally trained (we stopped quite early) but also b) the finetuning settings were chosen rather randomly. From our experience so far, you can almost always get performance increases for individual basins with finetuning, but it is difficult to find settings that are universally applicable. However, this tutorial was just a showcase of how easy it actually is to finetune models with the NeuralHydrology library. Now it is up to you to experiment with it.