Trainers API¶
Author: Wenyu Ouyang Date: 2021-12-05 11:21:58 LastEditTime: 2025-11-08 16:02:09 LastEditors: Wenyu Ouyang Description: Main function for training and testing FilePath: orchhydro orchhydro rainers rainer.py Copyright (c) 2021-2022 Wenyu Ouyang. All rights reserved.
ensemble_train_and_evaluate(cfgs)
¶
Function to train and test for ensemble models
Parameters¶
cfgs Dictionary containing all configs needed to run the model
Returns¶
None
Source code in torchhydro/trainers/trainer.py
def ensemble_train_and_evaluate(cfgs: Dict):
"""
Function to train and test for ensemble models
Parameters
----------
cfgs
Dictionary containing all configs needed to run the model
Returns
-------
None
"""
# for basins and models
ensemble = cfgs["training_cfgs"]["ensemble"]
if not ensemble:
raise ValueError(
"ensemble should be True, otherwise should use train_and_evaluate rather than ensemble_train_and_evaluate"
)
ensemble_items = cfgs["training_cfgs"]["ensemble_items"]
number_of_items = len(ensemble_items)
if number_of_items == 0:
raise ValueError("ensemble_items should not be empty")
keys_list = list(ensemble_items.keys())
if "kfold" in keys_list:
_trans_kfold_to_periods(cfgs, ensemble_items, "kfold")
_nested_loop_train_and_evaluate(keys_list, 0, ensemble_items, cfgs)
set_random_seed(seed)
¶
Set a random seed to guarantee the reproducibility
Parameters¶
seed a number
Returns¶
None
Source code in torchhydro/trainers/trainer.py
def set_random_seed(seed):
"""
Set a random seed to guarantee the reproducibility
Parameters
----------
seed
a number
Returns
-------
None
"""
# print("Random seed:", seed)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
train_and_evaluate(cfgs)
¶
Function to train and test a Model
Parameters¶
cfgs Dictionary containing all configs needed to run the model
Returns¶
None
Source code in torchhydro/trainers/trainer.py
def train_and_evaluate(cfgs: Dict):
"""
Function to train and test a Model
Parameters
----------
cfgs
Dictionary containing all configs needed to run the model
Returns
-------
None
"""
random_seed = cfgs["training_cfgs"]["random_seed"]
set_random_seed(random_seed)
resulter = Resulter(cfgs)
deephydro = _get_deep_hydro(cfgs)
# if train_mode is False, we only evaluate the model
train_mode = deephydro.cfgs["training_cfgs"]["train_mode"]
# but if train_mode is True, we still need some conditions to train the model
continue_train = deephydro.cfgs["model_cfgs"]["continue_train"]
is_transfer_learning = deephydro.cfgs["model_cfgs"]["model_type"] == "TransLearn"
is_train = train_mode and (
(deephydro.weight_path is not None and (continue_train or is_transfer_learning))
or (deephydro.weight_path is None)
)
if is_train:
deephydro.model_train()
# Explicitly delete optimizer to free its states (momentum, etc.)
"""if hasattr(deephydro, "optimizer"):
del deephydro.optimizer
# Note: deephydro.model will be re-initialized in model_evaluate
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
torch.mps.empty_cache()"""
preds, obss = deephydro.model_evaluate()
resulter.save_cfg(deephydro.cfgs)
resulter.save_result(preds, obss)
resulter.eval_result(preds, obss)
Author: Wenyu Ouyang Date: 2024-04-08 18:15:48 LastEditTime: 2025-07-13 16:25:31 LastEditors: Wenyu Ouyang Description: HydroDL model class FilePath: orchhydro orchhydro rainers\deep_hydro.py Copyright (c) 2024-2024 Wenyu Ouyang. All rights reserved.
DeepHydro (DeepHydroInterface)
¶
The Base Trainer class for Hydrological Deep Learning models
Source code in torchhydro/trainers/deep_hydro.py
class DeepHydro(DeepHydroInterface):
"""
The Base Trainer class for Hydrological Deep Learning models
"""
def __init__(
self,
cfgs: Dict,
pre_model=None,
):
"""
Parameters
----------
cfgs
configs for the model
pre_model
a pre-trained model, if it is not None,
we will use its weights to initialize this model
by default None
"""
super().__init__(cfgs)
self.split_validation = validate_split_ranges(self.cfgs)
# Initialize fabric based on configuration
self.fabric = create_fabric_wrapper(cfgs.get("training_cfgs", {}))
self.pre_model = pre_model
self.model = self.fabric.setup_module(self.load_model())
if cfgs["training_cfgs"]["train_mode"]:
self.traindataset = self.make_dataset("train")
self.traindataset.split_evaluation_semantics = self.split_validation[
"evaluation_semantics"
]
if cfgs["data_cfgs"]["t_range_valid"] is not None:
self.validdataset = self.make_dataset("valid")
self.validdataset.split_evaluation_semantics = self.split_validation[
"evaluation_semantics"
]
self.testdataset: BaseDataset = self.make_dataset("test")
self.testdataset.split_evaluation_semantics = self.split_validation[
"evaluation_semantics"
]
@property
def device(self):
"""Get the device from fabric wrapper"""
return self.fabric._device
def load_model(self, mode="train"):
"""
Load a time series forecast model in pytorch_model_dict in model_dict_function.py
Returns
-------
object
model in pytorch_model_dict in model_dict_function.py
"""
if mode == "infer":
if self.weight_path is None or self.cfgs["model_cfgs"]["continue_train"]:
# if no weight path is provided
# or weight file is provided but continue train again,
# we will use the trained model in the new case_dir directory
self.weight_path = self._get_trained_model()
elif mode != "train":
raise ValueError("Invalid mode; must be 'train' or 'infer'")
model_cfgs = self.cfgs["model_cfgs"]
model_name = model_cfgs["model_name"]
if model_name not in pytorch_model_dict:
raise NotImplementedError(
f"Error the model {model_name} was not found in the model dict. Please add it."
)
if self.pre_model is not None:
return self._load_pretrain_model()
elif self.weight_path is not None:
return self._load_model_from_pth()
else:
return pytorch_model_dict[model_name](**model_cfgs["model_hyperparam"])
def _load_pretrain_model(self):
"""load a pretrained model as the initial model"""
return self.pre_model
def _load_model_from_pth(self):
weight_path = self.weight_path
model_cfgs = self.cfgs["model_cfgs"]
model_name = model_cfgs["model_name"]
model = pytorch_model_dict[model_name](**model_cfgs["model_hyperparam"])
checkpoint = torch.load(
weight_path, map_location=self.device, weights_only=False
)
model.load_state_dict(checkpoint)
print("Weights sucessfully loaded")
return model
def make_dataset(self, is_tra_val_te: str):
"""
Initializes a pytorch dataset.
Parameters
----------
is_tra_val_te
train or valid or test
Returns
-------
object
an object initializing from class in datasets_dict in data_dict.py
"""
data_cfgs = self.cfgs["data_cfgs"]
dataset_name = data_cfgs["dataset"]
if dataset_name in list(datasets_dict.keys()):
dataset = datasets_dict[dataset_name](self.cfgs, is_tra_val_te)
else:
raise NotImplementedError(
f"Error the dataset {str(dataset_name)} was not found in the dataset dict. Please add it."
)
return dataset
def model_train(self) -> None:
"""train a hydrological DL model"""
# A dictionary of the necessary parameters for training
training_cfgs = self.cfgs["training_cfgs"]
# The file path to load model weights from; defaults to "model_save"
model_filepath = self.cfgs["data_cfgs"]["case_dir"]
data_cfgs = self.cfgs["data_cfgs"]
es = None
if training_cfgs["early_stopping"]:
es = EarlyStopper(training_cfgs["patience"])
criterion = self._get_loss_func(training_cfgs)
opt = self._get_optimizer(training_cfgs)
scheduler = self._get_scheduler(training_cfgs, opt)
max_epochs = training_cfgs["epochs"]
start_epoch = training_cfgs["start_epoch"]
# use PyTorch's DataLoader to load the data into batches in each epoch
data_loader, validation_data_loader = self._get_dataloader(
training_cfgs, data_cfgs
)
logger = TrainLogger(model_filepath, self.cfgs, opt)
performance_monitor = (
TrainingPerformanceMonitor(
Path(model_filepath) / "performance",
self.device,
gpu_sample_interval=float(
training_cfgs.get("performance_gpu_sample_interval", 1.0)
),
gpu_idle_threshold=float(
training_cfgs.get("performance_gpu_idle_threshold", 5.0)
),
)
if training_cfgs.get("performance_monitor", False)
else None
)
for epoch in range(start_epoch, max_epochs + 1):
if performance_monitor is not None:
performance_monitor.start_epoch(epoch)
with logger.log_epoch_train(epoch) as train_logs:
total_loss, n_iter_ep = torch_single_train(
self.model,
opt,
criterion,
data_loader,
device=self.device,
which_first_tensor=training_cfgs["which_first_tensor"],
non_blocking_transfer=training_cfgs.get(
"non_blocking_transfer",
bool(training_cfgs.get("pin_memory", False)),
),
cuda_prefetch=bool(training_cfgs.get("cuda_prefetch", False)),
cuda_prefetch_batches=int(
training_cfgs.get("cuda_prefetch_batches", 2)
),
performance_monitor=performance_monitor,
amp=bool(training_cfgs.get("amp", False)),
amp_dtype=training_cfgs.get("amp_dtype"),
mixed_precision=training_cfgs.get("mixed_precision", "off"),
nonfinite_check_interval=int(
training_cfgs.get("nonfinite_check_interval", 100)
),
)
train_logs["train_loss"] = total_loss
train_logs["model"] = self.model
if performance_monitor is not None:
performance = performance_monitor.finish_epoch(
sampler=getattr(data_loader, "batch_sampler", None),
dataset=self.traindataset,
)
print(
"Performance epoch {epoch}: {rate:.1f} samples/s, "
"loader wait {wait:.4f}s/batch, GPU util {gpu}, "
"GPU idle {idle}".format(
epoch=epoch,
rate=performance["samples_per_second"],
wait=performance["avg_dataloader_wait_seconds"] or 0.0,
gpu=(
f'{performance["gpu_util_avg_percent"]:.1f}%'
if performance["gpu_util_avg_percent"] is not None
else "n/a"
),
idle=(
f'{performance["gpu_idle_percent"]:.1f}%'
if performance["gpu_idle_percent"] is not None
else "n/a"
),
)
)
valid_loss = None
valid_metrics = None
if data_cfgs["t_range_valid"] is not None:
with logger.log_epoch_valid(epoch) as valid_logs:
valid_loss, valid_metrics = self._1epoch_valid(
training_cfgs, criterion, validation_data_loader, valid_logs
)
self._scheduler_step(training_cfgs, scheduler, valid_loss)
logger.save_session_param(
epoch, total_loss, n_iter_ep, valid_loss, valid_metrics
)
logger.save_model_and_params(self.model, epoch, self.cfgs)
if es and not es.check_loss(
self.model,
valid_loss,
self.cfgs["data_cfgs"]["case_dir"],
):
print("Stopping model now")
break
# logger.plot_model_structure(self.model)
logger.tb.close()
# return the trained model weights and bias and the epoch loss
return self.model.state_dict(), sum(logger.epoch_loss) / len(logger.epoch_loss)
def _get_scheduler(self, training_cfgs, opt):
lr_scheduler_cfg = training_cfgs["lr_scheduler"]
if "lr" in lr_scheduler_cfg and "lr_factor" not in lr_scheduler_cfg:
target_lr = lr_scheduler_cfg["lr"]
scheduler = LambdaLR(
opt, lr_lambda=lambda epoch: target_lr / opt.param_groups[0]["lr"]
)
# scheduler = LambdaLR(opt, lr_lambda=lambda epoch: 1.0)
elif isinstance(lr_scheduler_cfg, dict) and all(
isinstance(epoch, int) for epoch in lr_scheduler_cfg
):
# piecewise constant learning rate
epochs = sorted(lr_scheduler_cfg.keys())
values = [lr_scheduler_cfg[e] for e in epochs]
def lr_lambda(epoch):
idx = bisect.bisect_right(epochs, epoch) - 1
return 1.0 if idx < 0 else values[idx]
scheduler = LambdaLR(opt, lr_lambda=lr_lambda)
elif "lr_factor" in lr_scheduler_cfg and "lr_patience" not in lr_scheduler_cfg:
scheduler = ExponentialLR(opt, gamma=lr_scheduler_cfg["lr_factor"])
elif "lr_factor" in lr_scheduler_cfg:
scheduler = ReduceLROnPlateau(
opt,
mode="min",
factor=lr_scheduler_cfg["lr_factor"],
patience=lr_scheduler_cfg["lr_patience"],
)
else:
raise ValueError("Invalid lr_scheduler configuration")
return scheduler
def _scheduler_step(self, training_cfgs, scheduler, valid_loss):
lr_scheduler_cfg = training_cfgs["lr_scheduler"]
required_keys = {"lr_factor", "lr_patience"}
if required_keys.issubset(lr_scheduler_cfg.keys()):
scheduler.step(valid_loss)
else:
scheduler.step()
def _1epoch_valid(
self, training_cfgs, criterion, validation_data_loader, valid_logs
):
valid_obss_np, valid_preds_np, valid_loss = compute_validation(
self.model,
criterion,
validation_data_loader,
device=self.device,
which_first_tensor=training_cfgs["which_first_tensor"],
mixed_precision=training_cfgs.get("mixed_precision", "off"),
non_blocking=bool(
training_cfgs.get(
"non_blocking_transfer", training_cfgs.get("pin_memory", False)
)
),
empty_cache_during_validation=bool(
training_cfgs.get("empty_cache_during_validation", False)
),
)
valid_logs["valid_loss"] = valid_loss
if (
self.cfgs["training_cfgs"]["valid_batch_mode"] == "test"
and self.cfgs["training_cfgs"]["calc_metrics"]
):
# NOTE: Now we only evaluate the metrics for test-mode validation
target_col = self.cfgs["data_cfgs"]["target_cols"]
valid_metrics = evaluate_validation(
validation_data_loader,
valid_preds_np,
valid_obss_np,
self.cfgs["evaluation_cfgs"],
target_col,
)
valid_logs["valid_metrics"] = valid_metrics
return valid_loss, valid_metrics
return valid_loss, None
def _get_trained_model(self):
model_loader = self.cfgs["evaluation_cfgs"]["model_loader"]
model_pth_dir = self.cfgs["data_cfgs"]["case_dir"]
return read_pth_from_model_loader(model_loader, model_pth_dir)
def model_evaluate(self):
"""
Evaluate the model, and perform SHAP analysis if needed.
Returns
-------
preds_xr, obss_xr
"""
self.model = self.load_model(mode="infer").to(self.device)
data_cfgs = self.cfgs["data_cfgs"]
training_cfgs = self.cfgs["training_cfgs"]
# Create dataloader in advance as it is used for both inference and SHAP
test_dataloader = self._get_dataloader(training_cfgs, data_cfgs, mode="infer")
# Decide whether to perform SHAP analysis based on configuration
eval_cfgs = self.cfgs.get("evaluation_cfgs", {})
shap_cfgs = eval_cfgs.get("shap_cfgs", {})
# Inference returns X_all for SHAP analysis only if enabled
if shap_cfgs.get("enable", False):
pred_xr, obs_xr, X_all = self.inference(test_dataloader, return_inputs=True)
print("SHAP analysis enabled.")
else:
pred_xr, obs_xr = self.inference(test_dataloader, return_inputs=False)
X_all = None
print("SHAP analysis disabled.")
if shap_cfgs.get("enable", False):
self.run_shap_analysis(
X_all=X_all,
save_dir=shap_cfgs.get("save_dir", "./shap_results"),
max_sample=shap_cfgs.get("max_sample", 3000),
background_size=shap_cfgs.get("background_size", 100),
)
return pred_xr, obs_xr
def inference(self, dataloader, return_inputs=True):
self.model.eval()
preds = []
obss = []
X_all = []
with torch.no_grad():
for batch in dataloader:
# Handle Seq2Seq / multi-input format: batch[0] might be a
# list or tuple of inputs (e.g., [encoder, decoder] or (x, z)).
if isinstance(batch[0], (list, tuple)):
# Multi-input format: move each element to device and
# transpose time-series tensors to (seq, batch, feature).
x = [
inp.permute(1, 0, 2).to(self.device)
if inp.ndim == 3
else inp.to(self.device)
for inp in batch[0]
]
pred = self.model(*x)
# Transpose prediction back to (batch, seq, feature)
if isinstance(pred, torch.Tensor) and pred.ndim == 3:
pred = pred.permute(1, 0, 2)
else:
# Traditional tensor or separated static-input tuple.
x = _move_batch_value_to_device(
batch[0],
self.device,
self.cfgs["training_cfgs"].get("which_first_tensor")
== "sequence",
)
pred = self.model(x)
# --- TRANSPOSE BACK TO BATCH-FIRST ---
if (
self.cfgs["training_cfgs"].get("which_first_tensor")
== "sequence"
and isinstance(pred, torch.Tensor)
and pred.ndim == 3
):
pred = pred.transpose(0, 1)
y = batch[1].to(self.device).cpu()
# --- DO NOT TRANSPOSE LABELS ---
# Labels from _flood_event_collate_fn or default stack are already (Batch, Seq, Feat).
# Transposing them here would make them (Seq, Batch, Feat), which breaks concatenation.
# Handle models that return dict (e.g., Seq2Seq with attention)
if isinstance(pred, dict):
# If this was a sequence-first dict, permute its tensors
if (
self.cfgs["training_cfgs"].get("which_first_tensor")
== "sequence"
):
for k, v in pred.items():
if isinstance(v, torch.Tensor) and v.ndim == 3:
# Avoid double-transposition if it starts with 'f' (frequency branches)
# and we already know where it came from. Actually, safer to check shape.
# seq_first: (Seq, Batch, Feat). Batch size is usually smaller than Seq len.
# But in inference, Seq len (Flood events) can be smaller too.
# Let's assume if it came through model(x) it is Seq-first.
pred[k] = v.transpose(0, 1)
if "f1" in pred:
# Use f1 for 1-hour scale predictions
pred = pred["f1"]
# elif 'f0' in pred:
# Fallback to f0
# pred = pred['f0']
else:
# Generic handling for other dict outputs
pred_tensor = None
for key in ["predictions", "output", "pred", "logits"]:
if key in pred and isinstance(pred[key], torch.Tensor):
pred_tensor = pred[key]
break
if pred_tensor is None:
pred_tensor = next(iter(pred.values()))
pred = pred_tensor
preds.append(pred.cpu())
obss.append(y.cpu())
if return_inputs:
# Handle both formats for X_all
if isinstance(x, list):
# For Seq2Seq, store the encoder input (first element)
x_for_shap = x[0]
else:
x_for_shap = x
# Ensure SHAP input is Batch-first
if (
self.cfgs["training_cfgs"].get("which_first_tensor")
== "sequence"
and x_for_shap.ndim == 3
):
x_for_shap = x_for_shap.transpose(0, 1)
X_all.append(x_for_shap.cpu())
# Convert to numpy
# Pad sequence dimension independently for preds and obss
if len(preds) > 0:
max_p_s = max(p.shape[1] for p in preds)
max_o_s = max(o.shape[1] for o in obss)
padded_preds = []
padded_obss = []
for p, o in zip(preds, obss):
if p.shape[1] < max_p_s:
pad_p = torch.full(
(p.shape[0], max_p_s - p.shape[1], p.shape[2]), np.nan
)
p = torch.cat([p, pad_p], dim=1)
if o.shape[1] < max_o_s:
pad_o = torch.full(
(o.shape[0], max_o_s - o.shape[1], o.shape[2]), np.nan
)
o = torch.cat([o, pad_o], dim=1)
padded_preds.append(p)
padded_obss.append(o)
pred_np = torch.cat(padded_preds, 0).numpy()
obs_np = torch.cat(padded_obss, 0).numpy()
else:
pred_np = np.array([])
obs_np = np.array([])
# Convert to xarray ---- KEEP THIS (TorchHydro requirement)
obs_xr, pred_xr = get_preds_to_be_eval(
dataloader, self.cfgs["evaluation_cfgs"], pred_np, obs_np
)
if return_inputs:
return pred_xr, obs_xr, X_all
else:
return pred_xr, obs_xr
def run_shap_analysis(
self, X_all, save_dir, max_sample=3000, feature_names=None, background_size=100
):
import torch
import numpy as np
import shap
import matplotlib.pyplot as plt
import os
print("Running FAST SHAP analysis...")
eval_cfgs = self.cfgs.get("evaluation_cfgs", {}) or {}
shap_cfgs = eval_cfgs.get("shap_cfgs", {}) or {}
sequence_enable = bool(shap_cfgs.get("sequence_enable", False))
plot_mode = shap_cfgs.get("plot_mode")
if not isinstance(plot_mode, str):
plot_mode = "all"
plot_mode = plot_mode.lower()
plot_fast = True
plot_seq = True
plot_hybrid = True
plot_split = False
if plot_mode == "split_only":
plot_fast = False
plot_seq = False
plot_hybrid = False
plot_split = True
sequence_enable = True
# ======================================================
# 1. Optimize Memory & Diversity: Sample sparsely from MANY batches
# ======================================================
# X_all is a list of tensors, each (B, T, F)
if not X_all:
print("[WARNING] No data for SHAP analysis.")
return
# Get dimensions from first batch
_, T, F_total = X_all[0].shape
# We need max_sample samples.
# To ensure spatial diversity (static attributes variation), we must sample from many batches.
# Strategy: Randomly select N batches, then sample k items from each batch.
import random
n_batches = len(X_all)
# We want to sample from at least min(50, n_batches) batches to cover different basins
target_batch_count = min(50, n_batches)
# How many samples per batch?
# We want total ~ max_sample * 2 (safety margin)
samples_per_batch = max(1, (max_sample * 2) // target_batch_count)
print(
f"[INFO] Sampling strategy: {target_batch_count} batches, ~{samples_per_batch} samples/batch"
)
batch_indices = list(range(n_batches))
random.shuffle(batch_indices)
selected_batch_indices = batch_indices[:target_batch_count]
collected_subsets = []
for idx in selected_batch_indices:
batch = X_all[idx] # (B, T, F)
B = batch.shape[0]
# Flatten this batch to (B*T, F)
batch_flat = batch.reshape(-1, F_total)
M_batch = batch_flat.shape[0]
# Sample randomly from this batch
if M_batch > samples_per_batch:
# Random indices
indices = torch.randperm(M_batch)[:samples_per_batch]
batch_subset = batch_flat[indices]
else:
batch_subset = batch_flat
collected_subsets.append(batch_subset)
# Concatenate subsets
X_subset = torch.cat(collected_subsets, dim=0) # (N_total_subset, F)
# Handle NaNs: Replace NaNs with 0 (mean of normalized data)
if torch.isnan(X_subset).any():
print(
"[WARNING] Input data contains NaNs. Filling with 0 for SHAP analysis."
)
X_subset = torch.nan_to_num(X_subset, nan=0.0)
print(f"[INFO] Collected subset samples: {X_subset.shape[0]}")
# ======================================================
# 2. Flatten to 2D: Already flattened in loop
# ======================================================
X_flat = X_subset
M = X_flat.shape[0]
print(f"[INFO] Flattened subset samples: {M}")
# ======================================================
# 3. Sample input data for acceleration
# ======================================================
# Ensure we don't sample more than we have
sample_size = min(max_sample, M)
if sample_size < M:
# We already shuffled and sampled, but let's do one final random selection to be sure
idx = np.random.choice(M, sample_size, replace=False)
X_sample = X_flat[idx]
else:
X_sample = X_flat
X_sample_np = X_sample.cpu().numpy() # numpy，shape=(sample_size, F)
print(f"[INFO] Using {sample_size} samples for SHAP")
# ======================================================
# 4. Select background for SHAP
# ======================================================
# background_size is passed as argument
if M > background_size:
idx_bg = np.random.choice(M, background_size, replace=False)
background = X_flat[idx_bg]
else:
background = X_flat
background_np = background.cpu().numpy()
print(f"[INFO] Using {background.shape[0]} samples as SHAP background")
# ======================================================
# 5. Wrap the model: SamplingExplainer uses a python function
# ======================================================
# Move model to CPU to avoid MPS bugs with SHAP/LSTM
original_device = self.device
shap_device = torch.device("cpu")
print(
f"[INFO] Moving model to {shap_device} for SHAP analysis to avoid MPS errors..."
)
self.model.to(shap_device)
self.model.eval()
def model_predict(x_np):
"""
Args:
x_np: numpy array with shape (batch, F).
Returns:
numpy array with shape (batch, 1).
Note:
Converts numpy array to torch.Tensor, feeds it to the model,
and reshapes the output to (batch, 1).
"""
x_tensor = torch.tensor(
x_np, dtype=torch.float32, device=shap_device
) # shape=(B,F)
# SimpleLSTM (and default nn.LSTM) expects (Time, Batch, Features) if batch_first=False
# We treat SHAP samples as independent time steps (Time=1)
x_tensor = x_tensor.unsqueeze(0) # (1, B, F)
if torch.isnan(x_tensor).any():
print(f"[WARNING] SHAP input x_tensor contains NaNs!")
with torch.no_grad():
out = self.model(x_tensor) # (1, B, 1)
if torch.isnan(out).any():
print(f"[WARNING] SHAP model output contains NaNs!")
# Optional: Replace output NaNs with 0 to allow SHAP to continue (but this hides model issues)
# out = torch.nan_to_num(out, nan=0.0)
# --- SamplingExplainer requires 2D output ---
if out.shape[0] == 1:
out = out.squeeze(0)
else:
print(
f"[WARNING] Unexpected output shape {out.shape} in SHAP. Squeezing dim 0."
)
out = out.squeeze(0)
# Debug output range
# print(f"[DEBUG] SHAP model output range: [{out.min().item():.4f}, {out.max().item():.4f}]")
return out.cpu().numpy()
# ======================================================
# 6. Run SHAP
# ======================================================
print(f"[INFO] Computing SHAP values...")
explainer = shap.SamplingExplainer(model_predict, background_np)
# Increased nsamples for stability
shap_values = explainer.shap_values(X_sample_np, nsamples=200)
# Check shap_values
if isinstance(shap_values, list):
shap_values = shap_values[0] # Take the first output channel
print(f"[INFO] SHAP shape = {np.array(shap_values).shape}")
# Handle NaNs in SHAP values
if np.isnan(shap_values).any():
nan_count = np.isnan(shap_values).sum()
total_count = shap_values.size
print(
f"[WARNING] SHAP values contain {nan_count}/{total_count} NaNs! Filling with 0."
)
shap_values = np.nan_to_num(shap_values, nan=0.0)
# ======================================================
# 7. Construct feature names
# ======================================================
if feature_names is None:
# Attempt to retrieve real feature names from configuration and stat file
feature_names = self._get_feature_names(F_total)
if feature_names is not None:
print(f"[INFO] Using reconstructed real feature names: {feature_names}")
else:
print(
f"[WARNING] Could not reconstruct feature names. Using default names."
)
data_cfgs = self.cfgs.get("data_cfgs", {})
relevant_cols = data_cfgs.get("relevant_cols", [])
feature_names = [f"DYNAMIC_{i}" for i in range(len(relevant_cols))] + [
f"STATIC_{i}" for i in range(F_total - len(relevant_cols))
]
# ======================================================
# 8. Output plots
# ======================================================
os.makedirs(save_dir, exist_ok=True)
shap_values_fast = np.array(shap_values)
X_fast_np = X_sample_np
if plot_fast:
plt.figure()
shap.summary_plot(
shap_values,
X_sample_np,
feature_names=feature_names,
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_summary_fast.png", dpi=300, bbox_inches="tight"
)
plt.close()
plt.figure()
shap.summary_plot(
shap_values,
X_sample_np,
feature_names=feature_names,
plot_type="bar",
max_display=100,
show=False,
)
plt.savefig(f"{save_dir}/shap_bar_fast.png", dpi=300, bbox_inches="tight")
plt.close()
print(f"[INFO] SHAP analysis saved to {save_dir}")
if sequence_enable:
training_cfgs = self.cfgs.get("training_cfgs", {}) or {}
warmup_length = training_cfgs.get("warmup_length")
hindcast_length = training_cfgs.get("hindcast_length")
forecast_length = training_cfgs.get("forecast_length")
seq_len = shap_cfgs.get("sequence_length")
if not isinstance(seq_len, int):
if (
isinstance(warmup_length, int)
and isinstance(hindcast_length, int)
and isinstance(forecast_length, int)
):
seq_len = warmup_length + hindcast_length + forecast_length
else:
seq_len = T
seq_max_sample = shap_cfgs.get("sequence_max_sample")
if not isinstance(seq_max_sample, int):
seq_max_sample = min(128, max_sample)
seq_max_sample = max(1, seq_max_sample)
seq_background_size = shap_cfgs.get("sequence_background_size")
if not isinstance(seq_background_size, int):
seq_background_size = min(32, background_size)
seq_background_size = max(1, seq_background_size)
seq_nsamples = shap_cfgs.get("sequence_nsamples")
if not isinstance(seq_nsamples, int):
seq_nsamples = 50
time_agg = shap_cfgs.get("time_agg")
if not isinstance(time_agg, str):
time_agg = "abs_sum"
time_agg = time_agg.lower()
which_first_tensor = (
self.cfgs.get("training_cfgs", {}).get("which_first_tensor") or "batch"
)
import random
n_batches = len(X_all)
target_batch_count = min(50, n_batches)
seqs_per_batch = max(1, (seq_max_sample * 2) // max(1, target_batch_count))
batch_indices = list(range(n_batches))
random.shuffle(batch_indices)
selected_batch_indices = batch_indices[:target_batch_count]
seq_list = []
for idx in selected_batch_indices:
batch = X_all[idx] # (B, T, F)
if not isinstance(batch, torch.Tensor) or batch.ndim != 3:
continue
B = int(batch.shape[0])
if B <= 0:
continue
take = min(B, seqs_per_batch)
seq_indices = torch.randperm(B)[:take]
for sidx in seq_indices:
seq = batch[sidx] # (T, F)
if not isinstance(seq, torch.Tensor) or seq.ndim != 2:
continue
T_seq = int(seq.shape[0])
if T_seq <= 0:
continue
if T_seq >= seq_len:
seq_use = seq[-seq_len:]
else:
pad_len = seq_len - T_seq
pad = torch.zeros(
(pad_len, F_total), dtype=seq.dtype, device=seq.device
)
seq_use = torch.cat([pad, seq], dim=0)
seq_list.append(seq_use)
if not seq_list:
print(
"[WARNING] sequence_enable=True but no sequences collected; skipping."
)
else:
X_seq = torch.stack(seq_list, dim=0) # (N, seq_len, F)
if torch.isnan(X_seq).any():
X_seq = torch.nan_to_num(X_seq, nan=0.0)
if X_seq.shape[0] > seq_max_sample:
X_seq = X_seq[torch.randperm(X_seq.shape[0])[:seq_max_sample]]
N_seq = int(X_seq.shape[0])
bg_n = min(seq_background_size, N_seq)
if bg_n < N_seq:
idx_bg = np.random.choice(N_seq, bg_n, replace=False)
background_seq = X_seq[idx_bg]
else:
background_seq = X_seq
print(
f"[INFO] SHAP sequence enabled: N={N_seq}, T={seq_len}, F={F_total}, bg={int(background_seq.shape[0])}, nsamples={seq_nsamples}, time_agg={time_agg}"
)
self.model.to(original_device)
self.model.eval()
try:
from torch.nn.modules.rnn import RNNBase
except Exception:
RNNBase = ()
rnn_modules = []
rnn_training_states = []
for m in self.model.modules():
if isinstance(m, RNNBase):
rnn_modules.append(m)
rnn_training_states.append(m.training)
m.train()
class _ShapSeqWrapper(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
if which_first_tensor == "sequence":
x = x.transpose(0, 1)
out = self.model(x)
if isinstance(out, dict):
if "f1" in out:
out = out["f1"]
else:
out = next(iter(out.values()))
if isinstance(out, torch.Tensor) and out.ndim == 3:
out = (
out[-1]
if which_first_tensor == "sequence"
else out[:, -1]
)
if isinstance(out, torch.Tensor) and out.ndim == 1:
out = out.unsqueeze(1)
return out
wrapper = _ShapSeqWrapper(self.model).to(original_device)
background_t = background_seq.to(original_device).float()
sample_t = X_seq.to(original_device).float()
try:
with torch.backends.cudnn.flags(enabled=False):
explainer = shap.GradientExplainer(wrapper, background_t)
shap_values = explainer.shap_values(
sample_t, nsamples=seq_nsamples
)
finally:
for m, st in zip(rnn_modules, rnn_training_states):
m.train(st)
if isinstance(shap_values, list):
shap_values = shap_values[0]
shap_values = np.array(shap_values)
if shap_values.ndim == 4 and shap_values.shape[-1] == 1:
shap_values = shap_values[..., 0]
if np.isnan(shap_values).any():
shap_values = np.nan_to_num(shap_values, nan=0.0)
if time_agg == "sum":
shap_values_feat = shap_values.sum(axis=1)
else:
shap_values_feat = np.abs(shap_values).sum(axis=1)
X_feat = sample_t.detach().cpu().numpy().mean(axis=1)
suffix = "sum" if time_agg == "sum" else "abs_sum"
if plot_seq:
plt.figure()
shap.summary_plot(
shap_values_feat,
X_feat,
feature_names=feature_names,
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_summary_seq_{suffix}.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
plt.figure()
shap.summary_plot(
shap_values_feat,
X_feat,
feature_names=feature_names,
plot_type="bar",
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_bar_seq_{suffix}.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
print(f"[INFO] SHAP sequence plots saved to {save_dir}")
data_cfgs = self.cfgs.get("data_cfgs", {}) or {}
relevant_cols = data_cfgs.get("relevant_cols", []) or []
static_start = int(len(relevant_cols))
static_start = max(0, min(static_start, F_total))
hybrid_n = int(
min(512, shap_values_fast.shape[0], shap_values_feat.shape[0])
)
if hybrid_n > 0 and static_start < F_total:
if shap_values_fast.shape[0] > hybrid_n:
idx_fast = np.random.choice(
shap_values_fast.shape[0], hybrid_n, replace=False
)
else:
idx_fast = np.arange(shap_values_fast.shape[0])
if shap_values_feat.shape[0] > hybrid_n:
idx_seq = np.random.choice(
shap_values_feat.shape[0], hybrid_n, replace=False
)
else:
idx_seq = np.arange(shap_values_feat.shape[0])
hybrid_shap = shap_values_fast[idx_fast].copy()
hybrid_x = X_fast_np[idx_fast].copy()
x_seq_last = sample_t.detach().cpu().numpy()[:, -1, :]
seq_shap_sel = shap_values_feat[idx_seq]
seq_x_sel = x_seq_last[idx_seq]
hybrid_shap[:, static_start:] = seq_shap_sel[:, static_start:]
hybrid_x[:, static_start:] = seq_x_sel[:, static_start:]
hybrid_dynamic_mode = shap_cfgs.get("hybrid_dynamic_mode")
if not isinstance(hybrid_dynamic_mode, str):
hybrid_dynamic_mode = "fast"
hybrid_dynamic_mode = hybrid_dynamic_mode.lower()
dyn_seq_idx = []
if hybrid_dynamic_mode == "seq_all":
dyn_seq_idx = list(range(static_start))
elif hybrid_dynamic_mode == "seq_bucket0":
feature_buckets = (
(self.cfgs.get("model_cfgs", {}) or {})
.get("model_hyperparam", {})
.get("feature_buckets")
)
if (
isinstance(feature_buckets, (list, tuple))
and len(feature_buckets) >= static_start
):
dyn_seq_idx = [
i
for i in range(static_start)
if int(feature_buckets[i]) == 0
]
extra_seq_features = shap_cfgs.get("hybrid_dynamic_seq_features")
if (
isinstance(extra_seq_features, (list, tuple))
and static_start > 0
):
name_to_idx = {
str(n): i
for i, n in enumerate(feature_names[:static_start])
}
for fn in extra_seq_features:
idx = name_to_idx.get(str(fn))
if idx is not None:
dyn_seq_idx.append(idx)
if dyn_seq_idx:
dyn_seq_idx = sorted(set(dyn_seq_idx))
hybrid_shap[:, dyn_seq_idx] = seq_shap_sel[:, dyn_seq_idx]
hybrid_x[:, dyn_seq_idx] = seq_x_sel[:, dyn_seq_idx]
if plot_hybrid:
plt.figure()
shap.summary_plot(
hybrid_shap,
hybrid_x,
feature_names=feature_names,
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_summary_hybrid.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
plt.figure()
shap.summary_plot(
hybrid_shap,
hybrid_x,
feature_names=feature_names,
plot_type="bar",
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_bar_hybrid.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
print(f"[INFO] SHAP hybrid plots saved to {save_dir}")
if plot_split:
feature_buckets = (
(self.cfgs.get("model_cfgs", {}) or {})
.get("model_hyperparam", {})
.get("feature_buckets")
)
fb = None
if isinstance(feature_buckets, (list, tuple)):
if len(feature_buckets) >= F_total:
fb = [int(x) for x in feature_buckets[:F_total]]
elif len(feature_buckets) == static_start:
fb = [int(x) for x in feature_buckets] + [0] * (
F_total - static_start
)
if fb is None:
low_idx = list(range(static_start, F_total))
high_idx = list(range(0, static_start))
else:
low_idx = [i for i in range(F_total) if int(fb[i]) == 0]
high_idx = [i for i in range(F_total) if int(fb[i]) > 0]
if low_idx:
low_shap = hybrid_shap[:, low_idx]
low_x = hybrid_x[:, low_idx]
low_names = [feature_names[i] for i in low_idx]
plt.figure()
shap.summary_plot(
low_shap,
low_x,
feature_names=low_names,
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_summary_hybrid_lowfreq.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
plt.figure()
shap.summary_plot(
low_shap,
low_x,
feature_names=low_names,
plot_type="bar",
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_bar_hybrid_lowfreq.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
if high_idx:
high_shap = hybrid_shap[:, high_idx]
high_x = hybrid_x[:, high_idx]
high_names = [feature_names[i] for i in high_idx]
plt.figure()
shap.summary_plot(
high_shap,
high_x,
feature_names=high_names,
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_summary_hybrid_highfreq.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
plt.figure()
shap.summary_plot(
high_shap,
high_x,
feature_names=high_names,
plot_type="bar",
max_display=100,
show=False,
)
plt.savefig(
f"{save_dir}/shap_bar_hybrid_highfreq.png",
dpi=300,
bbox_inches="tight",
)
plt.close()
print(
f"[INFO] SHAP hybrid high/low frequency split plots saved to {save_dir}"
)
clean_extra = shap_cfgs.get("clean_extra")
if not isinstance(clean_extra, bool):
clean_extra = True
if clean_extra:
keep = {
"shap_bar_hybrid_lowfreq.png",
"shap_bar_hybrid_highfreq.png",
"shap_summary_hybrid_lowfreq.png",
"shap_summary_hybrid_highfreq.png",
}
for fn in os.listdir(save_dir):
if not (fn.startswith("shap_") and fn.endswith(".png")):
continue
if fn in keep:
continue
try:
os.remove(os.path.join(save_dir, fn))
except Exception:
pass
# Restore model to original device
self.model.to(original_device)
print(f"[INFO] Model restored to {original_device}")
def _get_feature_names(self, F_total):
"""
Reconstruct feature names, handling One-Hot Encoding and extra features (e.g. embeddings) found in stat_dict.
"""
import json
data_cfgs = self.cfgs.get("data_cfgs", {})
relevant_cols = data_cfgs.get("relevant_cols", [])
constant_cols = data_cfgs.get("constant_cols", [])
target_cols = data_cfgs.get("target_cols", [])
# Check for One-Hot Encoding
scaler_params = data_cfgs.get("scaler_params", {})
one_hot_cols = scaler_params.get("one_hot_cols", [])
# Try to load stat_dict
# Priority 1: Directory of the weight path (best_model.pth)
stat_dict = {}
stat_file_loaded = False
search_paths = []
if self.weight_path:
search_paths.append(
os.path.join(
os.path.dirname(self.weight_path), "dapengscaler_stat.json"
)
)
if "case_dir" in data_cfgs:
search_paths.append(
os.path.join(data_cfgs["case_dir"], "dapengscaler_stat.json")
)
for stat_file in search_paths:
if os.path.exists(stat_file):
try:
with open(stat_file, "r") as fp:
stat_dict = json.load(fp)
print(f"[INFO] Loaded stat file from {stat_file}")
stat_file_loaded = True
break
except Exception as e:
print(f"[WARNING] Failed to load stat file {stat_file}: {e}")
if not stat_file_loaded:
print(
"[WARNING] Could not load dapengscaler_stat.json from any expected location."
)
feature_names = []
# 1. Dynamic features (relevant_cols)
feature_names.extend(relevant_cols)
# 2. Static features (constant_cols) from config
# We need to track which keys we have used to identify "extra" ones later
used_keys = set(relevant_cols) | set(target_cols)
for col in constant_cols:
used_keys.add(col)
if col in one_hot_cols:
# Expand One-Hot columns
categories = stat_dict.get(f"{col}_categories", [])
if categories:
for cat in categories:
feature_names.append(f"{col}_{cat}")
else:
# Fallback
feature_names.append(f"{col}_UNKNOWN")
else:
feature_names.append(col)
# 3. Check for missing features (e.g. AlphaEarth embeddings)
if len(feature_names) < F_total and stat_dict:
print(
f"[INFO] Feature name count ({len(feature_names)}) < Model Input ({F_total}). Searching for extra features in stat_dict..."
)
extra_candidates = []
for key in stat_dict.keys():
# Skip known columns
if key in used_keys:
continue
# Skip One-Hot category metadata
if key.endswith("_categories"):
continue
# Skip One-Hot original columns if they were expanded (already handled by checking used_keys? No, one_hot_cols are in constant_cols, so in used_keys)
# Add to candidates
extra_candidates.append(key)
# Sort candidates to ensure deterministic order
# We assume extra features (like embeddings) are appended in alphabetical or numerical order
extra_candidates.sort()
# Try appending extra candidates
# We only append enough to match F_total if possible, or just all of them and see
# Heuristic: If we need exactly N more features, and we found M extra candidates.
needed = F_total - len(feature_names)
if len(extra_candidates) >= needed:
# If we have enough candidates, take them.
# If we have MORE than needed, it's ambiguous. But usually embeddings are distinct.
# Let's assume all valid extra candidates are part of the input.
# However, stat_dict might contain other auxiliary variables.
# If we match exactly, great.
if len(extra_candidates) == needed:
print(
f"[INFO] Found exactly {needed} extra features. Appending them."
)
feature_names.extend(extra_candidates)
else:
print(
f"[WARNING] Found {len(extra_candidates)} extra candidates, but needed {needed}. Appending first {needed} sorted keys."
)
feature_names.extend(extra_candidates[:needed])
else:
print(
f"[WARNING] Found {len(extra_candidates)} extra candidates, but needed {needed}. Appending all available."
)
feature_names.extend(extra_candidates)
# Final check
if len(feature_names) != F_total:
print(
f"[WARNING] Reconstructed feature names length ({len(feature_names)}) does not match model input ({F_total})."
)
return None
return feature_names
def _get_optimizer(self, training_cfgs):
params_in_opt = self.model.parameters()
return pytorch_opt_dict[training_cfgs["optimizer"]](
params_in_opt, **training_cfgs["optim_params"]
)
def _get_loss_func(self, training_cfgs):
criterion_init_params = {}
if "criterion_params" in training_cfgs:
loss_param = training_cfgs["criterion_params"]
if loss_param is not None:
for key in loss_param.keys():
if key in ("loss_funcs", "eval_loss_func"):
criterion_init_params[key] = pytorch_criterion_dict[
loss_param[key]
]()
else:
criterion_init_params[key] = loss_param[key]
return pytorch_criterion_dict[training_cfgs["criterion"]](
**criterion_init_params
)
def _flood_event_collate_fn(self, batch):
"""自定义的洪水事件 collate 函数，确保所有样本长度一致"""
# Check if this is Seq2Seq format: ([x1, x2, ...], y)
if isinstance(batch[0][0], list):
# Seq2Seq format: use default collate for list of tensors
from torch.utils.data.dataloader import default_collate
return default_collate(batch)
# 找到这个批次中最长的序列长度
max_len = max(tensor_data[0].shape[0] for tensor_data in batch)
# 调整所有样本到相同长度
processed_batch = []
for tensor_data in batch:
# 获取x和y（假设tensor_data[0]是x，tensor_data[1]是y）
x = tensor_data[0]
y = tensor_data[1] if len(tensor_data) > 1 else None
current_len = x.shape[0]
if current_len < max_len:
# 使用 NaN 填充 x
padding_x = torch.full(
(max_len - current_len, x.shape[1]), np.nan, dtype=x.dtype
)
padded_x = torch.cat([x, padding_x], dim=0)
# 如果有y，也进行填充
if y is not None:
padding_y = torch.full(
(max_len - current_len, y.shape[1]), np.nan, dtype=y.dtype
)
padded_y = torch.cat([y, padding_y], dim=0)
else:
padded_y = None
else:
# 如果更长则截断
padded_x = x[:max_len]
padded_y = y[:max_len] if y is not None else None
if padded_y is not None:
processed_batch.append((padded_x, padded_y))
else:
processed_batch.append(padded_x)
# 根据数据结构返回堆叠后的结果
if len(processed_batch) > 0 and isinstance(processed_batch[0], tuple):
return (
torch.stack([x for x, _ in processed_batch], 0),
torch.stack([y for _, y in processed_batch], 0),
)
else:
return torch.stack(processed_batch, 0)
def _get_dataloader(self, training_cfgs, data_cfgs, mode="train"):
if mode == "infer":
_collate_fn = None
# Use GNN collate function for GNN/graph datasets in inference mode
if hasattr(self.testdataset, "__class__") and (
"GNN" in self.testdataset.__class__.__name__
or "TgHydroDataset" in self.testdataset.__class__.__name__
):
_collate_fn = gnn_collate_fn
# 使用自定义的 collate 函数处理 FloodEventDataset
elif (
hasattr(self.testdataset, "__class__")
and "FloodEvent" in self.testdataset.__class__.__name__
):
_collate_fn = self._flood_event_collate_fn
return DataLoader(
self.testdataset,
batch_size=training_cfgs["batch_size"],
shuffle=False,
sampler=None,
batch_sampler=None,
drop_last=False,
timeout=0,
worker_init_fn=None,
collate_fn=_collate_fn,
)
worker_num = 0
pin_memory = False
if "num_workers" in training_cfgs:
worker_num = training_cfgs["num_workers"]
print(f"using {str(worker_num)} workers")
if "pin_memory" in training_cfgs:
pin_memory = training_cfgs["pin_memory"]
print(f"Pin memory set to {str(pin_memory)}")
dataloader_extra_kwargs = {}
if worker_num > 0:
if "persistent_workers" in training_cfgs:
dataloader_extra_kwargs["persistent_workers"] = training_cfgs[
"persistent_workers"
]
if "prefetch_factor" in training_cfgs:
dataloader_extra_kwargs["prefetch_factor"] = training_cfgs[
"prefetch_factor"
]
multiprocessing_context = training_cfgs.get("multiprocessing_context")
if multiprocessing_context not in (None, "auto"):
dataloader_extra_kwargs["multiprocessing_context"] = (
multiprocessing_context
)
if "in_order" in inspect.signature(DataLoader).parameters:
dataloader_extra_kwargs["in_order"] = bool(
training_cfgs.get("dataloader_in_order", False)
)
sampler = self._get_sampler(data_cfgs, training_cfgs, self.traindataset)
batch_sampler = sampler if getattr(sampler, "is_batch_sampler", False) else None
sampler_arg = None if batch_sampler is not None else sampler
_collate_fn = None
if training_cfgs["variable_length_cfgs"]["use_variable_length"]:
_collate_fn = varied_length_collate_fn
# Use GNN collate function for GNN/graph datasets
elif hasattr(self.traindataset, "__class__") and (
"GNN" in self.traindataset.__class__.__name__
or "TgHydroDataset" in self.traindataset.__class__.__name__
):
_collate_fn = gnn_collate_fn
use_batch_level_dataloader = bool(
training_cfgs.get("batch_level_dataloader", False)
)
separate_static_input = bool(
training_cfgs.get("separate_static_input", False)
)
if hasattr(self.traindataset, "_return_static_attrs_separately"):
self.traindataset._return_static_attrs_separately = False
if hasattr(self.traindataset, "_batch_level_dataloader_enabled"):
self.traindataset._batch_level_dataloader_enabled = False
if hasattr(self.traindataset, "_separate_static_input"):
self.traindataset._separate_static_input = False
train_supports_prebatched = bool(
getattr(self.traindataset, "supports_prebatched_getitems", False)
)
if (
(sampler is None or batch_sampler is not None)
and _collate_fn is None
and (
self.traindataset.__class__ is BaseDataset or train_supports_prebatched
)
):
self.traindataset._separate_static_input = separate_static_input
if (use_batch_level_dataloader or train_supports_prebatched) and hasattr(
self.traindataset, "_get_batch_level_items"
):
self.traindataset._batch_level_dataloader_enabled = True
_collate_fn = prebatched_collate_fn
elif self.traindataset.__class__ is BaseDataset:
self.traindataset._return_static_attrs_separately = True
_collate_fn = static_attr_batch_collate_fn
elif use_batch_level_dataloader:
warnings.warn(
"batch_level_dataloader=True was ignored because this "
"dataset uses a custom sampler/collate path or is not the "
"standard BaseDataset."
)
loader_kwargs = {
"num_workers": worker_num,
"pin_memory": pin_memory,
"timeout": (
float(training_cfgs.get("dataloader_timeout", 0.0))
if worker_num > 0
else 0
),
"collate_fn": _collate_fn,
**dataloader_extra_kwargs,
}
if batch_sampler is not None:
loader_kwargs["batch_sampler"] = batch_sampler
else:
loader_kwargs.update(
{
"batch_size": training_cfgs["batch_size"],
"shuffle": sampler_arg is None,
"sampler": sampler_arg,
}
)
data_loader = DataLoader(self.traindataset, **loader_kwargs)
if data_cfgs["t_range_valid"] is not None:
# Use the same collate function for validation dataset
_val_collate_fn = None
if training_cfgs["variable_length_cfgs"]["use_variable_length"]:
_val_collate_fn = varied_length_collate_fn
elif hasattr(self.validdataset, "__class__") and (
"GNN" in self.validdataset.__class__.__name__
or "TgHydroDataset" in self.validdataset.__class__.__name__
):
_val_collate_fn = gnn_collate_fn
if hasattr(self.validdataset, "_return_static_attrs_separately"):
self.validdataset._return_static_attrs_separately = False
if hasattr(self.validdataset, "_batch_level_dataloader_enabled"):
self.validdataset._batch_level_dataloader_enabled = False
if hasattr(self.validdataset, "_separate_static_input"):
self.validdataset._separate_static_input = False
valid_supports_prebatched = bool(
getattr(self.validdataset, "supports_prebatched_getitems", False)
)
if _val_collate_fn is None and (
self.validdataset.__class__ is BaseDataset or valid_supports_prebatched
):
self.validdataset._separate_static_input = separate_static_input
if (
use_batch_level_dataloader or valid_supports_prebatched
) and hasattr(self.validdataset, "_get_batch_level_items"):
self.validdataset._batch_level_dataloader_enabled = True
_val_collate_fn = prebatched_collate_fn
elif self.validdataset.__class__ is BaseDataset:
self.validdataset._return_static_attrs_separately = True
_val_collate_fn = static_attr_batch_collate_fn
elif use_batch_level_dataloader:
warnings.warn(
"batch_level_dataloader=True was ignored for validation "
"because this dataset uses a custom collate path or is "
"not the standard BaseDataset."
)
validation_data_loader = DataLoader(
self.validdataset,
batch_size=training_cfgs["batch_size"],
shuffle=False,
num_workers=worker_num,
pin_memory=pin_memory,
timeout=(
float(training_cfgs.get("dataloader_timeout", 0.0))
if worker_num > 0
else 0
),
collate_fn=_val_collate_fn,
**dataloader_extra_kwargs,
)
return data_loader, validation_data_loader
return data_loader, None
def _get_sampler(self, data_cfgs, training_cfgs, train_dataset):
"""
return data sampler based on the provided configuration and training dataset.
Parameters
----------
data_cfgs : dict
Configuration dictionary containing parameters for data sampling. Expected keys are:
- "sampler": dict, containing:
- "name": str, name of the sampler to use.
- "sampler_hyperparam": dict, optional hyperparameters for the sampler.
training_cfgs: dict
Configuration dictionary containing parameters for training. Expected keys are:
- "batch_size": int, size of each batch.
train_dataset : Dataset
The training dataset object which contains the data to be sampled. Expected attributes are:
- ngrid: int, number of grids in the dataset.
- nt: int, number of time steps in the dataset.
- rho: int, length of the input sequence.
- warmup_length: int, length of the warmup period.
- horizon: int, length of the forecast horizon.
Returns
-------
sampler_class
An instance of the specified sampler class, initialized with the provided dataset and hyperparameters.
Raises
------
NotImplementedError
If the specified sampler name is not found in the `data_sampler_dict`.
"""
if data_cfgs["sampler"] is None:
return None
batch_size = training_cfgs["batch_size"]
rho = train_dataset.rho
warmup_length = train_dataset.warmup_length
horizon = train_dataset.horizon
ngrid = train_dataset.ngrid
nt = train_dataset.nt
sampler_name = data_cfgs["sampler"]
if sampler_name not in data_sampler_dict:
raise NotImplementedError(f"Sampler {sampler_name} not implemented yet")
sampler_class = data_sampler_dict[sampler_name]
sampler_hyperparam = dict(data_cfgs.get("sampler_hyperparam") or {})
if sampler_name == "KuaiSampler":
sampler_hyperparam |= {
"batch_size": batch_size,
"warmup_length": warmup_length,
"rho_horizon": rho + horizon,
"ngrid": ngrid,
"nt": nt,
}
elif sampler_name == "WindowLenBatchSampler":
sampler_hyperparam |= {
"batch_size": batch_size,
}
elif sampler_name == "CloudZarrChunkBatchSampler":
sampler_hyperparam |= {
"batch_size": batch_size,
"worker_affinity_count": max(
1, int(training_cfgs.get("num_workers", 0))
),
}
return sampler_class(train_dataset, **sampler_hyperparam)
device
property
readonly
¶
Get the device from fabric wrapper
__init__(self, cfgs, pre_model=None)
special
¶
Parameters¶
cfgs configs for the model pre_model a pre-trained model, if it is not None, we will use its weights to initialize this model by default None
Source code in torchhydro/trainers/deep_hydro.py
def __init__(
self,
cfgs: Dict,
pre_model=None,
):
"""
Parameters
----------
cfgs
configs for the model
pre_model
a pre-trained model, if it is not None,
we will use its weights to initialize this model
by default None
"""
super().__init__(cfgs)
self.split_validation = validate_split_ranges(self.cfgs)
# Initialize fabric based on configuration
self.fabric = create_fabric_wrapper(cfgs.get("training_cfgs", {}))
self.pre_model = pre_model
self.model = self.fabric.setup_module(self.load_model())
if cfgs["training_cfgs"]["train_mode"]:
self.traindataset = self.make_dataset("train")
self.traindataset.split_evaluation_semantics = self.split_validation[
"evaluation_semantics"
]
if cfgs["data_cfgs"]["t_range_valid"] is not None:
self.validdataset = self.make_dataset("valid")
self.validdataset.split_evaluation_semantics = self.split_validation[
"evaluation_semantics"
]
self.testdataset: BaseDataset = self.make_dataset("test")
self.testdataset.split_evaluation_semantics = self.split_validation[
"evaluation_semantics"
]
load_model(self, mode='train')
¶
Load a time series forecast model in pytorch_model_dict in model_dict_function.py
Returns¶
object model in pytorch_model_dict in model_dict_function.py
Source code in torchhydro/trainers/deep_hydro.py
def load_model(self, mode="train"):
"""
Load a time series forecast model in pytorch_model_dict in model_dict_function.py
Returns
-------
object
model in pytorch_model_dict in model_dict_function.py
"""
if mode == "infer":
if self.weight_path is None or self.cfgs["model_cfgs"]["continue_train"]:
# if no weight path is provided
# or weight file is provided but continue train again,
# we will use the trained model in the new case_dir directory
self.weight_path = self._get_trained_model()
elif mode != "train":
raise ValueError("Invalid mode; must be 'train' or 'infer'")
model_cfgs = self.cfgs["model_cfgs"]
model_name = model_cfgs["model_name"]
if model_name not in pytorch_model_dict:
raise NotImplementedError(
f"Error the model {model_name} was not found in the model dict. Please add it."
)
if self.pre_model is not None:
return self._load_pretrain_model()
elif self.weight_path is not None:
return self._load_model_from_pth()
else:
return pytorch_model_dict[model_name](**model_cfgs["model_hyperparam"])
make_dataset(self, is_tra_val_te)
¶
Initializes a pytorch dataset.
Parameters¶
is_tra_val_te train or valid or test
Returns¶
object an object initializing from class in datasets_dict in data_dict.py
Source code in torchhydro/trainers/deep_hydro.py
def make_dataset(self, is_tra_val_te: str):
"""
Initializes a pytorch dataset.
Parameters
----------
is_tra_val_te
train or valid or test
Returns
-------
object
an object initializing from class in datasets_dict in data_dict.py
"""
data_cfgs = self.cfgs["data_cfgs"]
dataset_name = data_cfgs["dataset"]
if dataset_name in list(datasets_dict.keys()):
dataset = datasets_dict[dataset_name](self.cfgs, is_tra_val_te)
else:
raise NotImplementedError(
f"Error the dataset {str(dataset_name)} was not found in the dataset dict. Please add it."
)
return dataset
model_evaluate(self)
¶
Evaluate the model, and perform SHAP analysis if needed.
Returns¶
preds_xr, obss_xr
Source code in torchhydro/trainers/deep_hydro.py
def model_evaluate(self):
"""
Evaluate the model, and perform SHAP analysis if needed.
Returns
-------
preds_xr, obss_xr
"""
self.model = self.load_model(mode="infer").to(self.device)
data_cfgs = self.cfgs["data_cfgs"]
training_cfgs = self.cfgs["training_cfgs"]
# Create dataloader in advance as it is used for both inference and SHAP
test_dataloader = self._get_dataloader(training_cfgs, data_cfgs, mode="infer")
# Decide whether to perform SHAP analysis based on configuration
eval_cfgs = self.cfgs.get("evaluation_cfgs", {})
shap_cfgs = eval_cfgs.get("shap_cfgs", {})
# Inference returns X_all for SHAP analysis only if enabled
if shap_cfgs.get("enable", False):
pred_xr, obs_xr, X_all = self.inference(test_dataloader, return_inputs=True)
print("SHAP analysis enabled.")
else:
pred_xr, obs_xr = self.inference(test_dataloader, return_inputs=False)
X_all = None
print("SHAP analysis disabled.")
if shap_cfgs.get("enable", False):
self.run_shap_analysis(
X_all=X_all,
save_dir=shap_cfgs.get("save_dir", "./shap_results"),
max_sample=shap_cfgs.get("max_sample", 3000),
background_size=shap_cfgs.get("background_size", 100),
)
return pred_xr, obs_xr
model_train(self)
¶
train a hydrological DL model
Source code in torchhydro/trainers/deep_hydro.py
def model_train(self) -> None:
"""train a hydrological DL model"""
# A dictionary of the necessary parameters for training
training_cfgs = self.cfgs["training_cfgs"]
# The file path to load model weights from; defaults to "model_save"
model_filepath = self.cfgs["data_cfgs"]["case_dir"]
data_cfgs = self.cfgs["data_cfgs"]
es = None
if training_cfgs["early_stopping"]:
es = EarlyStopper(training_cfgs["patience"])
criterion = self._get_loss_func(training_cfgs)
opt = self._get_optimizer(training_cfgs)
scheduler = self._get_scheduler(training_cfgs, opt)
max_epochs = training_cfgs["epochs"]
start_epoch = training_cfgs["start_epoch"]
# use PyTorch's DataLoader to load the data into batches in each epoch
data_loader, validation_data_loader = self._get_dataloader(
training_cfgs, data_cfgs
)
logger = TrainLogger(model_filepath, self.cfgs, opt)
performance_monitor = (
TrainingPerformanceMonitor(
Path(model_filepath) / "performance",
self.device,
gpu_sample_interval=float(
training_cfgs.get("performance_gpu_sample_interval", 1.0)
),
gpu_idle_threshold=float(
training_cfgs.get("performance_gpu_idle_threshold", 5.0)
),
)
if training_cfgs.get("performance_monitor", False)
else None
)
for epoch in range(start_epoch, max_epochs + 1):
if performance_monitor is not None:
performance_monitor.start_epoch(epoch)
with logger.log_epoch_train(epoch) as train_logs:
total_loss, n_iter_ep = torch_single_train(
self.model,
opt,
criterion,
data_loader,
device=self.device,
which_first_tensor=training_cfgs["which_first_tensor"],
non_blocking_transfer=training_cfgs.get(
"non_blocking_transfer",
bool(training_cfgs.get("pin_memory", False)),
),
cuda_prefetch=bool(training_cfgs.get("cuda_prefetch", False)),
cuda_prefetch_batches=int(
training_cfgs.get("cuda_prefetch_batches", 2)
),
performance_monitor=performance_monitor,
amp=bool(training_cfgs.get("amp", False)),
amp_dtype=training_cfgs.get("amp_dtype"),
mixed_precision=training_cfgs.get("mixed_precision", "off"),
nonfinite_check_interval=int(
training_cfgs.get("nonfinite_check_interval", 100)
),
)
train_logs["train_loss"] = total_loss
train_logs["model"] = self.model
if performance_monitor is not None:
performance = performance_monitor.finish_epoch(
sampler=getattr(data_loader, "batch_sampler", None),
dataset=self.traindataset,
)
print(
"Performance epoch {epoch}: {rate:.1f} samples/s, "
"loader wait {wait:.4f}s/batch, GPU util {gpu}, "
"GPU idle {idle}".format(
epoch=epoch,
rate=performance["samples_per_second"],
wait=performance["avg_dataloader_wait_seconds"] or 0.0,
gpu=(
f'{performance["gpu_util_avg_percent"]:.1f}%'
if performance["gpu_util_avg_percent"] is not None
else "n/a"
),
idle=(
f'{performance["gpu_idle_percent"]:.1f}%'
if performance["gpu_idle_percent"] is not None
else "n/a"
),
)
)
valid_loss = None
valid_metrics = None
if data_cfgs["t_range_valid"] is not None:
with logger.log_epoch_valid(epoch) as valid_logs:
valid_loss, valid_metrics = self._1epoch_valid(
training_cfgs, criterion, validation_data_loader, valid_logs
)
self._scheduler_step(training_cfgs, scheduler, valid_loss)
logger.save_session_param(
epoch, total_loss, n_iter_ep, valid_loss, valid_metrics
)
logger.save_model_and_params(self.model, epoch, self.cfgs)
if es and not es.check_loss(
self.model,
valid_loss,
self.cfgs["data_cfgs"]["case_dir"],
):
print("Stopping model now")
break
# logger.plot_model_structure(self.model)
logger.tb.close()
# return the trained model weights and bias and the epoch loss
return self.model.state_dict(), sum(logger.epoch_loss) / len(logger.epoch_loss)
DeepHydroInterface (ABC)
¶
An abstract class used to handle different configurations of hydrological deep learning models + hyperparams for training, test, and predict functions. This class assumes that data is already split into test train and validation at this point.
Source code in torchhydro/trainers/deep_hydro.py
class DeepHydroInterface(ABC):
"""
An abstract class used to handle different configurations
of hydrological deep learning models + hyperparams for training, test, and predict functions.
This class assumes that data is already split into test train and validation at this point.
"""
def __init__(self, cfgs: Dict):
"""
Parameters
----------
cfgs
configs for initializing DeepHydro
"""
self._cfgs = cfgs
@property
def cfgs(self):
"""all configs"""
return self._cfgs
@property
def weight_path(self):
"""weight path"""
return self._cfgs["model_cfgs"]["weight_path"]
@weight_path.setter
def weight_path(self, weight_path):
self._cfgs["model_cfgs"]["weight_path"] = weight_path
@abstractmethod
def load_model(self, mode="train") -> object:
"""Get a Hydro DL model"""
raise NotImplementedError
@abstractmethod
def make_dataset(self, is_tra_val_te: str) -> object:
"""
Initializes a pytorch dataset.
Parameters
----------
is_tra_val_te
train or valid or test
Returns
-------
object
a dataset class loading data from data source
"""
raise NotImplementedError
@abstractmethod
def model_train(self):
"""
Train the model
"""
raise NotImplementedError
@abstractmethod
def model_evaluate(self):
"""
Evaluate the model
"""
raise NotImplementedError
cfgs
property
readonly
¶
all configs
weight_path
property
writable
¶
weight path
__init__(self, cfgs)
special
¶
Parameters¶
cfgs configs for initializing DeepHydro
Source code in torchhydro/trainers/deep_hydro.py
def __init__(self, cfgs: Dict):
"""
Parameters
----------
cfgs
configs for initializing DeepHydro
"""
self._cfgs = cfgs
load_model(self, mode='train')
¶
Get a Hydro DL model
Source code in torchhydro/trainers/deep_hydro.py
@abstractmethod
def load_model(self, mode="train") -> object:
"""Get a Hydro DL model"""
raise NotImplementedError
make_dataset(self, is_tra_val_te)
¶
Initializes a pytorch dataset.
Parameters¶
is_tra_val_te train or valid or test
Returns¶
object a dataset class loading data from data source
Source code in torchhydro/trainers/deep_hydro.py
@abstractmethod
def make_dataset(self, is_tra_val_te: str) -> object:
"""
Initializes a pytorch dataset.
Parameters
----------
is_tra_val_te
train or valid or test
Returns
-------
object
a dataset class loading data from data source
"""
raise NotImplementedError
model_evaluate(self)
¶
Evaluate the model
Source code in torchhydro/trainers/deep_hydro.py
@abstractmethod
def model_evaluate(self):
"""
Evaluate the model
"""
raise NotImplementedError
model_train(self)
¶
Train the model
Source code in torchhydro/trainers/deep_hydro.py
@abstractmethod
def model_train(self):
"""
Train the model
"""
raise NotImplementedError
FedLearnHydro (DeepHydro)
¶
Federated Learning Hydrological DL model
Source code in torchhydro/trainers/deep_hydro.py
class FedLearnHydro(DeepHydro):
"""Federated Learning Hydrological DL model"""
def __init__(self, cfgs: Dict):
super().__init__(cfgs)
# a user group which is a dict where the keys are the user index
# and the values are the corresponding data for each of those users
train_dataset = self.traindataset
fl_hyperparam = self.cfgs["model_cfgs"]["fl_hyperparam"]
# sample training data amongst users
if fl_hyperparam["fl_sample"] == "basin":
# Sample a basin for a user
user_groups = fl_sample_basin(train_dataset)
elif fl_hyperparam["fl_sample"] == "region":
# Sample a region for a user
user_groups = fl_sample_region(train_dataset)
else:
raise NotImplementedError()
self.user_groups = user_groups
@property
def num_users(self):
"""number of users in federated learning"""
return len(self.user_groups)
def model_train(self) -> None:
# BUILD MODEL
global_model = self.model
# copy weights
global_weights = global_model.state_dict()
# Training
train_loss, train_accuracy = [], []
print_every = 2
training_cfgs = self.cfgs["training_cfgs"]
model_cfgs = self.cfgs["model_cfgs"]
max_epochs = training_cfgs["epochs"]
start_epoch = training_cfgs["start_epoch"]
fl_hyperparam = model_cfgs["fl_hyperparam"]
# total rounds in a FL system is max_epochs
for epoch in tqdm(range(start_epoch, max_epochs + 1)):
local_weights, local_losses = [], []
print(f"\n | Global Training Round : {epoch} |\n")
global_model.train()
m = max(int(fl_hyperparam["fl_frac"] * self.num_users), 1)
# randomly select m users, they will be the clients in this round
idxs_users = np.random.choice(range(self.num_users), m, replace=False)
for idx in idxs_users:
# each user will be used to train the model locally
# user_gourps[idx] means the idx of dataset for a user
user_cfgs = self._get_a_user_cfgs(idx)
local_model = DeepHydro(
user_cfgs,
pre_model=copy.deepcopy(global_model),
)
w, loss = local_model.model_train()
local_weights.append(copy.deepcopy(w))
local_losses.append(copy.deepcopy(loss))
# update global weights
global_weights = average_weights(local_weights)
# update global weights
global_model.load_state_dict(global_weights)
loss_avg = sum(local_losses) / len(local_losses)
train_loss.append(loss_avg)
# Calculate avg training accuracy over all users at every epoch
list_acc = []
global_model.eval()
for c in range(self.num_users):
one_user_cfg = self._get_a_user_cfgs(c)
local_model = DeepHydro(
one_user_cfg,
pre_model=global_model,
)
acc, _, _ = local_model.model_evaluate()
list_acc.append(acc)
values = [list(d.values())[0][0] for d in list_acc]
filtered_values = [v for v in values if not np.isnan(v)]
train_accuracy.append(sum(filtered_values) / len(filtered_values))
# print global training loss after every 'i' rounds
if (epoch + 1) % print_every == 0:
print(f" \nAvg Training Stats after {epoch+1} global rounds:")
print(f"Training Loss : {np.mean(np.array(train_loss))}")
print("Train Accuracy: {:.2f}% \n".format(100 * train_accuracy[-1]))
def _get_a_user_cfgs(self, idx):
"""To get a user's configs for local training"""
user = self.user_groups[idx]
# update data_cfgs
# Use defaultdict to collect dates for each basin
basin_dates = defaultdict(list)
for _, (basin, time) in user.items():
basin_dates[basin].append(time)
# Initialize a list to store distinct basins
basins = []
# for each basin, we can find its date range
date_ranges = {}
for basin, times in basin_dates.items():
basins.append(basin)
date_ranges[basin] = (np.min(times), np.max(times))
# get the longest date range
longest_date_range = max(date_ranges.values(), key=lambda x: x[1] - x[0])
# transform the date range of numpy data into string
longest_date_range = [
np.datetime_as_string(dt, unit="D") for dt in longest_date_range
]
user_cfgs = copy.deepcopy(self.cfgs)
# update data_cfgs
update_nested_dict(
user_cfgs, ["data_cfgs", "t_range_train"], longest_date_range
)
# for local training in FL, we don't need a validation set
update_nested_dict(user_cfgs, ["data_cfgs", "t_range_valid"], None)
# for local training in FL, we don't need a test set, but we should set one to avoid error
update_nested_dict(user_cfgs, ["data_cfgs", "t_range_test"], longest_date_range)
update_nested_dict(user_cfgs, ["data_cfgs", "object_ids"], basins)
# update training_cfgs
# we also need to update some training params for local training from FL settings
update_nested_dict(
user_cfgs,
["training_cfgs", "epochs"],
user_cfgs["model_cfgs"]["fl_hyperparam"]["fl_local_ep"],
)
update_nested_dict(
user_cfgs,
["evaluation_cfgs", "test_epoch"],
user_cfgs["model_cfgs"]["fl_hyperparam"]["fl_local_ep"],
)
# don't need to save model weights for local training
update_nested_dict(
user_cfgs,
["training_cfgs", "save_epoch"],
None,
)
# there are two settings for batch size in configs, we need to update both of them
update_nested_dict(
user_cfgs,
["training_cfgs", "batch_size"],
user_cfgs["model_cfgs"]["fl_hyperparam"]["fl_local_bs"],
)
update_nested_dict(
user_cfgs,
["data_cfgs", "batch_size"],
user_cfgs["model_cfgs"]["fl_hyperparam"]["fl_local_bs"],
)
# update model_cfgs finally
# For local model, its model_type is Normal
update_nested_dict(user_cfgs, ["model_cfgs", "model_type"], "Normal")
update_nested_dict(
user_cfgs,
["model_cfgs", "fl_hyperparam"],
None,
)
return user_cfgs
num_users
property
readonly
¶
number of users in federated learning
model_train(self)
¶
train a hydrological DL model
Source code in torchhydro/trainers/deep_hydro.py
def model_train(self) -> None:
# BUILD MODEL
global_model = self.model
# copy weights
global_weights = global_model.state_dict()
# Training
train_loss, train_accuracy = [], []
print_every = 2
training_cfgs = self.cfgs["training_cfgs"]
model_cfgs = self.cfgs["model_cfgs"]
max_epochs = training_cfgs["epochs"]
start_epoch = training_cfgs["start_epoch"]
fl_hyperparam = model_cfgs["fl_hyperparam"]
# total rounds in a FL system is max_epochs
for epoch in tqdm(range(start_epoch, max_epochs + 1)):
local_weights, local_losses = [], []
print(f"\n | Global Training Round : {epoch} |\n")
global_model.train()
m = max(int(fl_hyperparam["fl_frac"] * self.num_users), 1)
# randomly select m users, they will be the clients in this round
idxs_users = np.random.choice(range(self.num_users), m, replace=False)
for idx in idxs_users:
# each user will be used to train the model locally
# user_gourps[idx] means the idx of dataset for a user
user_cfgs = self._get_a_user_cfgs(idx)
local_model = DeepHydro(
user_cfgs,
pre_model=copy.deepcopy(global_model),
)
w, loss = local_model.model_train()
local_weights.append(copy.deepcopy(w))
local_losses.append(copy.deepcopy(loss))
# update global weights
global_weights = average_weights(local_weights)
# update global weights
global_model.load_state_dict(global_weights)
loss_avg = sum(local_losses) / len(local_losses)
train_loss.append(loss_avg)
# Calculate avg training accuracy over all users at every epoch
list_acc = []
global_model.eval()
for c in range(self.num_users):
one_user_cfg = self._get_a_user_cfgs(c)
local_model = DeepHydro(
one_user_cfg,
pre_model=global_model,
)
acc, _, _ = local_model.model_evaluate()
list_acc.append(acc)
values = [list(d.values())[0][0] for d in list_acc]
filtered_values = [v for v in values if not np.isnan(v)]
train_accuracy.append(sum(filtered_values) / len(filtered_values))
# print global training loss after every 'i' rounds
if (epoch + 1) % print_every == 0:
print(f" \nAvg Training Stats after {epoch+1} global rounds:")
print(f"Training Loss : {np.mean(np.array(train_loss))}")
print("Train Accuracy: {:.2f}% \n".format(100 * train_accuracy[-1]))
TransLearnHydro (DeepHydro)
¶
Source code in torchhydro/trainers/deep_hydro.py
class TransLearnHydro(DeepHydro):
def __init__(self, cfgs: Dict, pre_model=None):
super().__init__(cfgs, pre_model)
def load_model(self, mode="train"):
"""Load model for transfer learning"""
model_cfgs = self.cfgs["model_cfgs"]
if self.weight_path is None and self.pre_model is None:
raise NotImplementedError(
"For transfer learning, we need a pre-trained model"
)
if mode == "train":
model = super().load_model(mode)
elif mode == "infer":
self.weight_path = self._get_trained_model()
model = self._load_model_from_pth()
model.to(self.device)
if (
"weight_path_add" in model_cfgs
and "freeze_params" in model_cfgs["weight_path_add"]
):
freeze_params = model_cfgs["weight_path_add"]["freeze_params"]
for param in freeze_params:
exec(f"model.{param}.requires_grad = False")
return model
def _load_model_from_pth(self):
weight_path = self.weight_path
model_cfgs = self.cfgs["model_cfgs"]
model_name = model_cfgs["model_name"]
model = pytorch_model_dict[model_name](**model_cfgs["model_hyperparam"])
checkpoint = torch.load(
weight_path, map_location=self.device, weights_only=False
)
if "weight_path_add" in model_cfgs:
if "excluded_layers" in model_cfgs["weight_path_add"]:
# delete some layers from source model if we don't need them
excluded_layers = model_cfgs["weight_path_add"]["excluded_layers"]
for layer in excluded_layers:
del checkpoint[layer]
print("sucessfully deleted layers")
else:
print("directly loading identically-named layers of source model")
model.load_state_dict(checkpoint, strict=False)
print("Weights sucessfully loaded")
return model
load_model(self, mode='train')
¶
Load model for transfer learning
Source code in torchhydro/trainers/deep_hydro.py
def load_model(self, mode="train"):
"""Load model for transfer learning"""
model_cfgs = self.cfgs["model_cfgs"]
if self.weight_path is None and self.pre_model is None:
raise NotImplementedError(
"For transfer learning, we need a pre-trained model"
)
if mode == "train":
model = super().load_model(mode)
elif mode == "infer":
self.weight_path = self._get_trained_model()
model = self._load_model_from_pth()
model.to(self.device)
if (
"weight_path_add" in model_cfgs
and "freeze_params" in model_cfgs["weight_path_add"]
):
freeze_params = model_cfgs["weight_path_add"]["freeze_params"]
for param in freeze_params:
exec(f"model.{param}.requires_grad = False")
return model
Author: Wenyu Ouyang Date: 2024-04-08 18:16:26 LastEditTime: 2025-12-07 09:07:22 LastEditors: Wenyu Ouyang Description: Some basic functions for training FilePath: orchhydro orchhydro rainers rain_utils.py Copyright (c) 2024-2024 Wenyu Ouyang. All rights reserved.
EarlyStopper
¶
Source code in torchhydro/trainers/train_utils.py
class EarlyStopper(object):
def __init__(
self,
patience: int,
min_delta: float = 0.0,
cumulative_delta: bool = False,
):
"""
EarlyStopping handler can be used to stop the training if no improvement after a given number of events.
Parameters
----------
patience
Number of events to wait if no improvement and then stop the training.
min_delta
A minimum increase in the score to qualify as an improvement,
i.e. an increase of less than or equal to `min_delta`, will count as no improvement.
cumulative_delta
It True, `min_delta` defines an increase since the last `patience` reset, otherwise,
it defines an increase after the last event. Default value is False.
"""
if patience < 1:
raise ValueError("Argument patience should be positive integer.")
if min_delta < 0.0:
raise ValueError("Argument min_delta should not be a negative number.")
self.patience = patience
self.min_delta = min_delta
self.cumulative_delta = cumulative_delta
self.counter = 0
self.best_score = None
def check_loss(self, model, validation_loss, save_dir) -> bool:
score = validation_loss
if self.best_score is None:
self.save_model_checkpoint(model, save_dir)
self.best_score = score
elif score + self.min_delta >= self.best_score:
self.counter += 1
print("Epochs without Model Update:", self.counter)
if self.counter >= self.patience:
return False
else:
self.save_model_checkpoint(model, save_dir)
print("Model Update")
self.best_score = score
self.counter = 0
return True
def save_model_checkpoint(self, model, save_dir):
torch.save(model.state_dict(), os.path.join(save_dir, "best_model.pth"))
__init__(self, patience, min_delta=0.0, cumulative_delta=False)
special
¶
EarlyStopping handler can be used to stop the training if no improvement after a given number of events.
Parameters¶
patience
Number of events to wait if no improvement and then stop the training.
min_delta
A minimum increase in the score to qualify as an improvement,
i.e. an increase of less than or equal to min_delta, will count as no improvement.
cumulative_delta
It True, min_delta defines an increase since the last patience reset, otherwise,
it defines an increase after the last event. Default value is False.
Source code in torchhydro/trainers/train_utils.py
def __init__(
self,
patience: int,
min_delta: float = 0.0,
cumulative_delta: bool = False,
):
"""
EarlyStopping handler can be used to stop the training if no improvement after a given number of events.
Parameters
----------
patience
Number of events to wait if no improvement and then stop the training.
min_delta
A minimum increase in the score to qualify as an improvement,
i.e. an increase of less than or equal to `min_delta`, will count as no improvement.
cumulative_delta
It True, `min_delta` defines an increase since the last `patience` reset, otherwise,
it defines an increase after the last event. Default value is False.
"""
if patience < 1:
raise ValueError("Argument patience should be positive integer.")
if min_delta < 0.0:
raise ValueError("Argument min_delta should not be a negative number.")
self.patience = patience
self.min_delta = min_delta
self.cumulative_delta = cumulative_delta
self.counter = 0
self.best_score = None
average_weights(w)
¶
Returns the average of the weights.
Source code in torchhydro/trainers/train_utils.py
def average_weights(w):
"""
Returns the average of the weights.
"""
w_avg = copy.deepcopy(w[0])
for key in w_avg.keys():
for i in range(1, len(w)):
w_avg[key] += w[i][key]
w_avg[key] = torch.div(w_avg[key], len(w))
return w_avg
cellstates_when_inference(seq_first, data_cfgs, pred)
¶
get cell states when inference
Source code in torchhydro/trainers/train_utils.py
def cellstates_when_inference(seq_first, data_cfgs, pred):
"""get cell states when inference"""
cs_out = (
pred.detach().cpu().numpy().swapaxes(0, 1)
if seq_first
else pred.detach().cpu().numpy()
)
cs_out_lst = [cs_out]
cell_state = reduce(lambda a, b: np.vstack((a, b)), cs_out_lst)
np.save(os.path.join(data_cfgs["case_dir"], "cell_states.npy"), cell_state)
# model.zero_grad()
torch.cuda.empty_cache()
return pred, cell_state
compute_loss(labels, output, criterion, **kwargs)
¶
Function for computing the loss
Parameters¶
labels The real values for the target. Shape can be variable but should follow (batch_size, time) output The output of the model. Can be a tensor or a dict with 'loss' key for generative models. criterion loss function validation_dataset Only passed when unscaling of data is needed. m defaults to 1
Returns¶
torch.Tensor the computed loss
Source code in torchhydro/trainers/train_utils.py
def compute_loss(labels: torch.Tensor, output, criterion, **kwargs) -> torch.Tensor:
"""
Function for computing the loss
Parameters
----------
labels
The real values for the target. Shape can be variable but should follow (batch_size, time)
output
The output of the model. Can be a tensor or a dict with 'loss' key for generative models.
criterion
loss function
validation_dataset
Only passed when unscaling of data is needed.
m
defaults to 1
Returns
-------
torch.Tensor
the computed loss
"""
# Handle generative models that compute loss internally
# These models return a dict with 'loss' key
if isinstance(output, dict) and "loss" in output:
return output["loss"]
if isinstance(criterion, GaussianLoss):
if len(output[0].shape) > 2:
g_loss = GaussianLoss(output[0][:, :, 0], output[1][:, :, 0])
else:
g_loss = GaussianLoss(output[0][:, 0], output[1][:, 0])
return g_loss(labels)
if isinstance(criterion, FloodBaseLoss):
# labels has one more column than output, which is the flood mask
# so we need to remove the last column of labels to get targets
flood_mask = labels[:, :, -1:] # Extract flood mask from last column
targets = labels[:, :, :-1] # Extract targets (remove last column)
return criterion(output, targets, flood_mask)
if isinstance(criterion, PyGraphLoss):
node_mask = kwargs.get("node_mask", None)
return criterion(output, labels, node_mask=node_mask)
# MultiOutLoss: 适配 PyG 批次的 [batch*num_nodes, horizon] 到 [time, batch, num_nodes]
if isinstance(criterion, MultiOutLoss):
batch_vector = kwargs.get("batch_vector", None)
node_mask = kwargs.get("node_mask", None)
if output.ndim == 2 and labels.ndim == 2 and batch_vector is not None:
# Output/Labels: [B*N_total, H]
# PyG batch processing: Directly compute loss on flattened data for efficiency
# 1. Apply Node Mask
if node_mask is not None:
if not isinstance(node_mask, torch.Tensor):
node_mask = torch.tensor(
node_mask, dtype=torch.bool, device=labels.device
)
else:
node_mask = node_mask.to(labels.device).bool()
# Ensure mask is 1D [B*N_total]
if node_mask.dim() > 1:
node_mask = node_mask.flatten()
output = output[node_mask]
labels = labels[node_mask]
# 2. Handle NaN in labels (consistent with MultiOutLoss behavior)
# Mask out any NaN values in the target
mask_nan = ~torch.isnan(labels)
output = output[mask_nan]
labels = labels[mask_nan]
# 3. Compute Loss
# Assume homogeneous loss function for all nodes in PyG batch
loss_func = criterion.loss_funcs
if isinstance(loss_func, list):
# Use the first loss function if a list is provided
loss_func = loss_func[0]
return loss_func(output, labels)
elif output.ndim == 2 and labels.ndim == 2:
# 无 batch_vector 情况，退化为 batch=1
output = output.transpose(0, 1).unsqueeze(1)
labels = labels.transpose(0, 1).unsqueeze(1)
n_out = labels.shape[-1]
try:
iw = criterion.item_weight
if not isinstance(iw, list) or len(iw) != n_out:
criterion.item_weight = [1.0] * n_out
except Exception:
pass
try:
dg = criterion.data_gap
if isinstance(dg, list) and len(dg) != n_out:
criterion.data_gap = [0] * n_out
except Exception:
pass
return criterion(output, labels.float())
if (
isinstance(output, torch.Tensor)
and len(labels.shape) != len(output.shape)
and len(labels.shape) > 1
):
if labels.shape[1] == output.shape[1]:
labels = labels.unsqueeze(2)
else:
labels = labels.unsqueeze(0)
assert labels.shape == output.shape
valid_target = torch.isfinite(labels)
if not valid_target.all():
if not valid_target.any():
return output.sum() * 0.0
output = output[valid_target]
labels = labels[valid_target]
return criterion(output, labels.float())
compute_validation(model, criterion, data_loader, device=None, **kwargs)
¶
Function to compute the validation loss metrics
Parameters¶
model the trained model criterion torch.nn.modules.loss dataloader The data-loader of either validation or test-data device torch.device
Returns¶
tuple validation observations (numpy array), predictions (numpy array) and the loss of validation
Source code in torchhydro/trainers/train_utils.py
def compute_validation(
model,
criterion,
data_loader: DataLoader,
device: torch.device = None,
**kwargs,
):
"""
Function to compute the validation loss metrics
Parameters
----------
model
the trained model
criterion
torch.nn.modules.loss
dataloader
The data-loader of either validation or test-data
device
torch.device
Returns
-------
tuple
validation observations (numpy array), predictions (numpy array) and the loss of validation
"""
model.eval()
seq_first = kwargs["which_first_tensor"] != "batch"
mixed_precision = kwargs.get("mixed_precision", "off")
non_blocking = bool(kwargs.get("non_blocking", False))
empty_cache_during_validation = bool(
kwargs.get(
"empty_cache_during_validation",
getattr(data_loader.dataset, "training_cfgs", {}).get(
"empty_cache_during_validation", False
),
)
)
valid_loss = 0.0
obs_parts = []
pred_parts = []
with torch.no_grad():
iter_num = 0
for batch in tqdm(data_loader, desc="Evaluating", total=len(data_loader)):
batch_vector = getattr(batch, "batch", None)
node_mask = getattr(batch, "node_mask", None)
with mixed_precision_autocast(device, mixed_precision):
trg, output = model_infer(
seq_first, device, model, batch, non_blocking=non_blocking
)
valid_loss_ = compute_loss(
trg,
output,
criterion,
batch_vector=batch_vector,
node_mask=node_mask,
)
if torch.isnan(valid_loss_):
# for not-train mode, we may get all nan data for trg
# so we skip this batch
continue
print("NAN loss detected, skipping this batch")
valid_loss = valid_loss + valid_loss_.item()
iter_num = iter_num + 1
# For flood datasets, remove the flood_mask column from observations
# to match the prediction dimensions for evaluation
trg_for_eval = (
trg[:, :, :-1] if isinstance(criterion, FloodBaseLoss) else trg
)
obs_parts.append(trg_for_eval.detach().float().cpu())
pred_parts.append(output.detach().float().cpu())
del trg, output
if empty_cache_during_validation and torch.cuda.is_available():
torch.cuda.empty_cache()
valid_loss = valid_loss / iter_num
obs_final = torch.cat(obs_parts, dim=0)
pred_final = torch.cat(pred_parts, dim=0)
y_obs = obs_final.numpy()
y_pred = pred_final.numpy()
# Adaptation for PyGraphLoss: Convert Time-Major to Basin-Major
if isinstance(criterion, PyGraphLoss):
dataset = data_loader.dataset
# Assuming the dataset has 'sites_id' in 't_s_dict' which indicates number of nodes/basins
if hasattr(dataset, "t_s_dict") and "sites_id" in dataset.t_s_dict:
num_nodes = len(dataset.t_s_dict["sites_id"])
if y_obs.shape[0] % num_nodes == 0:
time_steps = y_obs.shape[0] // num_nodes
# Reshape from [Time*Nodes, ...] to [Time, Nodes, ...]
# Then transpose to [Nodes, Time, ...]
# Then flatten to [Nodes*Time, ...]
# Handle arbitrary trailing dimensions
# Ensure we have feature dimension. If shape is [Batch, Horizon], add feature dim -> [Batch, Horizon, 1]
if y_obs.ndim == 2:
y_obs = y_obs[..., np.newaxis]
if y_pred.ndim == 2:
y_pred = y_pred[..., np.newaxis]
obs_shape = y_obs.shape
pred_shape = y_pred.shape
# Reshape to [Time, Nodes, Horizon, Features]
y_obs = y_obs.reshape(time_steps, num_nodes, *obs_shape[1:])
# Transpose to [Nodes, Time, Horizon, Features]
permute_dims = [1, 0] + list(range(2, y_obs.ndim))
y_obs = y_obs.transpose(permute_dims)
# Same for predictions
y_pred = y_pred.reshape(time_steps, num_nodes, *pred_shape[1:])
y_pred = y_pred.transpose(permute_dims)
return y_obs, y_pred, valid_loss
create_amp_grad_scaler(device, mixed_precision)
¶
Create a GradScaler for CUDA FP16 training; BF16 does not need scaling.
Source code in torchhydro/trainers/train_utils.py
def create_amp_grad_scaler(device: torch.device | None, mixed_precision):
"""Create a GradScaler for CUDA FP16 training; BF16 does not need scaling."""
mode = _normalize_mixed_precision_mode(mixed_precision)
enabled = bool(mode == "fp16" and device is not None and device.type == "cuda")
if hasattr(torch, "amp") and hasattr(torch.amp, "GradScaler"):
try:
return torch.amp.GradScaler("cuda", enabled=enabled)
except TypeError:
return torch.amp.GradScaler(enabled=enabled)
return torch.cuda.amp.GradScaler(enabled=enabled)
evaluate_validation(validation_data_loader, output, labels, evaluation_cfgs, target_col)
¶
calculate metrics for validation
Parameters¶
output model output labels model target evaluation_cfgs evaluation configs target_col target columns
Returns¶
tuple metrics
Source code in torchhydro/trainers/train_utils.py
def evaluate_validation(
validation_data_loader,
output,
labels,
evaluation_cfgs,
target_col,
):
"""
calculate metrics for validation
Parameters
----------
output
model output
labels
model target
evaluation_cfgs
evaluation configs
target_col
target columns
Returns
-------
tuple
metrics
"""
fill_nan = evaluation_cfgs["fill_nan"]
if isinstance(fill_nan, list) and len(fill_nan) != len(target_col):
raise ValueError("Length of fill_nan must be equal to length of target_col.")
eval_log = {}
evaluation_metrics = evaluation_cfgs["metrics"]
obss_xr, preds_xr = get_preds_to_be_eval(
validation_data_loader,
evaluation_cfgs,
output,
labels,
)
# obss_xr_list
# preds_xr_list
# if type()
# for i in range(obs.shape[0]): # 第几个预见期
## obs_ = obs[i]
if isinstance(obss_xr, list):
obss_xr_list = obss_xr
preds_xr_list = preds_xr
for horizon_idx in range(len(obss_xr_list)):
obss_xr = obss_xr_list[horizon_idx]
preds_xr = preds_xr_list[horizon_idx]
for i, col in enumerate(target_col):
obs = obss_xr[col].to_numpy()
pred = preds_xr[col].to_numpy()
# eval_log will be updated rather than completely replaced, no need to use eval_log["key"]
eval_log = calculate_and_record_metrics(
obs,
pred,
evaluation_metrics,
col,
fill_nan[i] if isinstance(fill_nan, list) else fill_nan,
eval_log,
horizon_idx + 1,
)
return eval_log
for i, col in enumerate(target_col):
obs = obss_xr[col].to_numpy()
pred = preds_xr[col].to_numpy()
# eval_log will be updated rather than completely replaced, no need to use eval_log["key"]
eval_log = calculate_and_record_metrics(
obs,
pred,
evaluation_metrics,
col,
fill_nan[i] if isinstance(fill_nan, list) else fill_nan,
eval_log,
)
return eval_log
get_latest_pbm_param_file(param_dir)
¶
Get the latest parameter file of physics-based models in the current directory.
Parameters¶
param_dir : str The directory of parameter files.
Returns¶
str The latest parameter file.
Source code in torchhydro/trainers/train_utils.py
def get_latest_pbm_param_file(param_dir):
"""Get the latest parameter file of physics-based models in the current directory.
Parameters
----------
param_dir : str
The directory of parameter files.
Returns
-------
str
The latest parameter file.
"""
param_file_lst = [
os.path.join(param_dir, f)
for f in os.listdir(param_dir)
if f.startswith("pb_params") and f.endswith(".csv")
]
param_files = [Path(f) for f in param_file_lst]
param_file_names_lst = [param_file.stem.split("_") for param_file in param_files]
ctimes = [
int(param_file_names[param_file_names.index("params") + 1])
for param_file_names in param_file_names_lst
]
return param_files[ctimes.index(max(ctimes))] if ctimes else None
get_latest_tensorboard_event_file(log_dir)
¶
Get the latest event file in the log_dir directory.
Parameters¶
log_dir : str The directory where the event files are stored.
Returns¶
str The latest event file.
Source code in torchhydro/trainers/train_utils.py
def get_latest_tensorboard_event_file(log_dir):
"""Get the latest event file in the log_dir directory.
Parameters
----------
log_dir : str
The directory where the event files are stored.
Returns
-------
str
The latest event file.
"""
event_file_lst = [
os.path.join(log_dir, f) for f in os.listdir(log_dir) if f.startswith("events")
]
event_files = [Path(f) for f in event_file_lst]
event_file_names_lst = [event_file.stem.split(".") for event_file in event_files]
ctimes = [
int(event_file_names[event_file_names.index("tfevents") + 1])
for event_file_names in event_file_names_lst
]
return event_files[ctimes.index(max(ctimes))]
get_masked_tensors(variable_length_cfgs, batch, seq_first)
¶
Get the mask for the data
Parameters¶
variable_length_cfgs : dict The variable length configuration batch : tuple or list or torch_geometric.data.Batch The batch data from collate_fn or dataset (can be PyG Batch object) seq_first : bool Whether the data is in sequence first format
Returns¶
tuple For standard datasets: (xs, ys, xs_mask, ys_mask, xs_lens, ys_lens) For GNN datasets: (xs, ys, edge_index, edge_weight, xs_mask, ys_mask, xs_lens, ys_lens) For GNN with batch vector: (xs, ys, edge_index, edge_weight, batch_vector, xs_mask, ys_mask, xs_lens, ys_lens)
Source code in torchhydro/trainers/train_utils.py
def get_masked_tensors(variable_length_cfgs, batch, seq_first):
"""Get the mask for the data
Parameters
----------
variable_length_cfgs : dict
The variable length configuration
batch : tuple or list or torch_geometric.data.Batch
The batch data from collate_fn or dataset (can be PyG Batch object)
seq_first : bool
Whether the data is in sequence first format
Returns
-------
tuple
For standard datasets: (xs, ys, xs_mask, ys_mask, xs_lens, ys_lens)
For GNN datasets: (xs, ys, edge_index, edge_weight, xs_mask, ys_mask, xs_lens, ys_lens)
For GNN with batch vector: (xs, ys, edge_index, edge_weight, batch_vector, xs_mask, ys_mask, xs_lens, ys_lens)
"""
# Check if batch is a PyG Batch/Data object
try:
from torch_geometric.data import Batch as PyGBatch, Data as PyGData
is_pyg_batch = isinstance(batch, (PyGBatch, PyGData))
except ImportError:
is_pyg_batch = False
xs_mask = None
ys_mask = None
xs_lens = None
ys_lens = None
edge_index = None
edge_weight = None
batch_vector = None
# Handle PyG Data/Batch objects
if is_pyg_batch:
# Extract data from PyG Batch
xs = (
batch.x
) # Node features: [batch_nodes, hindcast_length, features] or similar
ys = batch.y # Target values: [batch_nodes, forecast_length] or similar
edge_index = batch.edge_index # Edge connectivity: [2, num_edges]
edge_weight = batch.edge_attr if hasattr(batch, "edge_attr") else None
batch_vector = batch.batch if hasattr(batch, "batch") else None
# For PyG batch, return appropriate format
if batch_vector is not None:
return (
xs,
ys,
edge_index,
edge_weight,
batch_vector,
xs_mask,
ys_mask,
xs_lens,
ys_lens,
)
elif edge_index is not None:
return xs, ys, edge_index, edge_weight, xs_mask, ys_mask, xs_lens, ys_lens
else:
return xs, ys, xs_mask, ys_mask, xs_lens, ys_lens
if variable_length_cfgs is None:
# Check batch length to determine format
if len(batch) == 5:
# GNN batch with batch_vector: [sxc, y, edge_index, edge_weight, batch_vector]
xs, ys, edge_index, edge_weight, batch_vector = batch
return (
xs,
ys,
edge_index,
edge_weight,
batch_vector,
xs_mask,
ys_mask,
xs_lens,
ys_lens,
)
elif len(batch) == 4:
# GNN batch: [sxc, y, edge_index, edge_weight]
xs, ys, edge_index, edge_weight = batch
return xs, ys, edge_index, edge_weight, xs_mask, ys_mask, xs_lens, ys_lens
else:
# Standard batch: [xs, ys]
xs, ys = batch[0], batch[1]
return xs, ys, xs_mask, ys_mask, xs_lens, ys_lens
if variable_length_cfgs.get("use_variable_length", False):
# When using variable length training, batch comes from varied_length_collate_fn
# which returns [xs_pad, ys_pad, xs_lens, ys_lens, xs_mask, ys_mask]
if len(batch) >= 6:
xs, ys, xs_lens, ys_lens, xs_mask_bool, ys_mask_bool = batch[:6]
else:
# Fallback: treat as regular batch with first two elements
xs, ys = batch[0], batch[1]
xs_lens = ys_lens = xs_mask_bool = ys_mask_bool = None
if xs_mask_bool is None and ys_mask_bool is None:
# sometime even you choose to use variable length training, the batch data may still be fixed length
# so we need to return the batch data directly
return xs, ys, xs_mask_bool, ys_mask_bool, xs_lens, ys_lens
# Convert masks to the format expected by model (float tensor with shape [..., 1])
xs_mask = xs_mask_bool.unsqueeze(-1).float() # [batch, seq, 1]
ys_mask = ys_mask_bool.unsqueeze(-1).float() # [batch, seq, 1]
# Convert to appropriate format for model if needed
if seq_first:
xs_mask = xs_mask.transpose(0, 1) # [seq, batch, 1]
ys_mask = ys_mask.transpose(0, 1) # [seq, batch, 1]
else:
# Check batch length to determine format
if len(batch) == 5:
# GNN batch with batch_vector: [sxc, y, edge_index, edge_weight, batch_vector]
xs, ys, edge_index, edge_weight, batch_vector = batch
elif len(batch) == 4:
# GNN batch: [sxc, y, edge_index, edge_weight]
xs, ys, edge_index, edge_weight = batch
else:
# Standard batch: [xs, ys]
xs, ys = batch[0], batch[1]
# Return appropriate format based on what we have
if edge_index is not None and edge_weight is not None:
return (
xs,
ys,
edge_index,
edge_weight,
batch_vector,
xs_mask,
ys_mask,
xs_lens,
ys_lens,
)
else:
return xs, ys, xs_mask, ys_mask, xs_lens, ys_lens
get_mixed_precision_dtype(device, mixed_precision)
¶
Return autocast dtype for the requested mixed precision mode.
Source code in torchhydro/trainers/train_utils.py
def get_mixed_precision_dtype(device: torch.device | None, mixed_precision) -> torch.dtype | None:
"""Return autocast dtype for the requested mixed precision mode."""
mode = _normalize_mixed_precision_mode(mixed_precision)
if mode == "off":
return None
if device is None:
raise ValueError("mixed precision requires an explicit torch.device.")
if device.type not in {"cuda", "cpu"}:
raise ValueError(
f"mixed precision mode {mode!r} is not supported on device {device}."
)
if mode == "bf16":
if device.type == "cuda" and not torch.cuda.is_bf16_supported():
raise RuntimeError("CUDA device does not report BF16 support.")
return torch.bfloat16
return torch.float16
get_preds_to_be_eval(valorte_data_loader, evaluation_cfgs, output, labels)
¶
Get prediction results prepared for evaluation: the denormalized data without metrics by different eval ways
Parameters¶
valorte_data_loader : DataLoader validation or test data loader evaluation_cfgs : dict evaluation configs output : np.ndarray model output labels : np.ndarray model target
Returns¶
tuple description
Source code in torchhydro/trainers/train_utils.py
def get_preds_to_be_eval(
valorte_data_loader,
evaluation_cfgs,
output,
labels,
):
"""
Get prediction results prepared for evaluation:
the denormalized data without metrics by different eval ways
Parameters
----------
valorte_data_loader : DataLoader
validation or test data loader
evaluation_cfgs : dict
evaluation configs
output : np.ndarray
model output
labels : np.ndarray
model target
Returns
-------
tuple
_description_
"""
evaluator = evaluation_cfgs["evaluator"]
# this test_rolling means how we perform prediction during testing
test_rolling = evaluation_cfgs["rolling"]
batch_size = valorte_data_loader.batch_size
target_scaler = valorte_data_loader.dataset.target_scaler
target_data = target_scaler.data_target
rho = valorte_data_loader.dataset.rho
horizon = valorte_data_loader.dataset.horizon
warmup_length = valorte_data_loader.dataset.warmup_length
hindcast_output_window = target_scaler.data_cfgs["hindcast_output_window"]
nf = valorte_data_loader.dataset.noutputvar # number of features
# number of time steps after warmup as outputs typically don't include warmup period
nt = valorte_data_loader.dataset.nt - warmup_length
basin_num = len(target_data.basin)
data_shape = (basin_num, nt, nf)
# Check if GNN mode (PyG output)
is_gnn = False
dataset = valorte_data_loader.dataset
if hasattr(dataset, "lookup_table") and len(dataset.lookup_table) > 0:
# lookup_table[0] is (basin, time, len). For GNN, basin is None.
if dataset.lookup_table[0][0] is None:
is_gnn = True
if evaluator["eval_way"] == "once":
if is_gnn:
raise NotImplementedError(
"eval_way='once' is not currently supported for PyG/GNN models. "
"Please use eval_way='rolling' which supports automatic reshaping of GNN outputs."
)
stride = evaluator["stride"]
if stride > 0:
if horizon != stride:
raise NotImplementedError(
"horizon should be equal to stride in evaluator if you chose eval_way to be once, or else you need to change the eval_way to be 1pace or rolling"
)
obs = _rolling_preds_for_once_eval(
(basin_num, horizon, nf),
rho,
evaluation_cfgs["forecast_length"],
stride,
hindcast_output_window,
target_data.reshape(basin_num, horizon, nf),
)
pred = _rolling_preds_for_once_eval(
(basin_num, horizon, nf),
rho,
evaluation_cfgs["forecast_length"],
stride,
hindcast_output_window,
output.reshape(batch_size, horizon, nf),
)
else:
if test_rolling > 0:
raise RuntimeError(
"please set rolling to 0 when you chose eval way as once and stride=0"
)
obs = labels.reshape(basin_num, -1, nf)
pred = output.reshape(basin_num, -1, nf)
elif evaluator["eval_way"] == "1pace":
if test_rolling < 1:
raise NotImplementedError(
"rolling should be larger than 0 if you chose eval_way to be 1pace"
)
pace_idx = evaluator["pace_idx"]
# stride = evaluator.get("stride", 1)
# for 1pace with pace_idx meaning which value of output was chosen to show
# 1st, we need to transpose data to 4-dim to show the whole data
# TODO:check should we select which def
pred = _recover_samples_to_basin(output, valorte_data_loader, pace_idx)
obs = _recover_samples_to_basin(labels, valorte_data_loader, pace_idx)
elif evaluator["eval_way"] == "rolling":
# 获取滚动预测所需的参数
stride = evaluator.get("stride", 1)
if stride != 1:
raise NotImplementedError(
"if stride is not equal to 1, we think it is meaningless"
)
# 重组预测结果和观测值
basin_num = len(target_data.basin)
# 新增：根据配置选择不同的数据组织方式
recover_mode = evaluator.get("recover_mode", "bybasins")
stride = evaluator.get("stride", 1)
data_shape = (basin_num, nt, nf)
if recover_mode == "bybasins":
pred = _recover_samples_to_4d_by_basins(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
output,
)
obs = _recover_samples_to_4d_by_basins(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
labels,
)
elif recover_mode == "byforecast":
pred = _recover_samples_to_4d_by_forecast(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
output, # samples, seq_length, nf
)
obs = _recover_samples_to_4d_by_forecast(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
labels,
)
elif recover_mode == "byensembles":
pred = _recover_samples_to_3d_by_4d_ensembles(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
output,
)
obs = _recover_samples_to_3d_by_4d_ensembles(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
labels,
)
else:
raise ValueError(
f"Unsupported recover_mode: {recover_mode}, must be 'bybasins' or 'byforecast' or 'byensembles'"
)
elif evaluator["eval_way"] == "floodevent":
# For flood event evaluation, stride is not typically used, but we set it to 1 for consistency
stride = evaluator.get("stride", 1)
pred = _recover_samples_to_continuous_by_floodevent(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
output,
)
obs = _recover_samples_to_continuous_by_floodevent(
data_shape,
valorte_data_loader,
stride,
hindcast_output_window,
labels,
)
else:
raise ValueError("eval_way should be rolling or 1pace")
# pace_idx = np.nan
recover_mode = evaluator.get("recover_mode")
valte_dataset = valorte_data_loader.dataset
# 检查数据维度并进行适当处理
if pred.ndim == 4:
# 如果是四维数据，需要根据评估方式选择合适的处理方法
if evaluator["eval_way"] == "1pace" and "pace_idx" in evaluator:
# 对于1pace模式，选择特定的预测步长
pace_idx = evaluator["pace_idx"]
# 选择特定预测步长的数据
pred_3d = pred[:, :, pace_idx, :]
obs_3d = obs[:, :, pace_idx, :]
preds_xr = valte_dataset.denormalize(pred_3d, pace_idx)
obss_xr = valte_dataset.denormalize(obs_3d, pace_idx)
elif evaluator["eval_way"] == "rolling" and recover_mode == "byforecast":
# 对于byforecast模式，需要特殊处理
# 创建一个列表存储每个预测步长的结果
preds_xr_list = []
obss_xr_list = []
if is_gnn:
start_idx = rho + 1
else:
start_idx = rho
for i in range(pred.shape[2]):
pred_3d = pred[:, :, i, :]
obs_3d = obs[:, :, i, :]
the_array_pred_ = np.full(target_data.shape, np.nan)
the_array_obs_ = np.full(target_data.shape, np.nan)
start = start_idx + i # TODO:need check
end = start + pred_3d.shape[1]
assert end <= the_array_pred_.shape[1]
the_array_pred_[:, start:end, :] = pred_3d
the_array_obs_[:, start:end, :] = obs_3d
preds_xr_list.append(valte_dataset.denormalize(the_array_pred_, i))
obss_xr_list.append(valte_dataset.denormalize(the_array_obs_, i))
# 合并结果
# preds_xr = xr.concat(preds_xr_list, dim="horizon")
# obss_xr = xr.concat(obss_xr_list, dim="horizon")
return obss_xr_list, preds_xr_list
elif evaluator["eval_way"] == "rolling" and recover_mode == "bybasins":
# 对于其他情况，可以考虑将四维数据转换为三维
# 例如，取最后一个预测步长
# 每个流域取最后一个预测步长 j = forecast_length-1,其绝对时间为
# warmup + τ + rho + (forecast_length-1),与 byforecast 分支对
# 最后一个步长(start_idx = rho + i)的放置一致。
selected_data = target_scaler.data_target
the_array_pred_ = np.full(selected_data.shape, np.nan)
the_array_obs_ = np.full(selected_data.shape, np.nan)
start = rho + pred.shape[2] - 1 # pred.shape[2] = forecast_length
assert (
start + pred.shape[1] <= the_array_pred_.shape[1]
), "填充范围超出目标数组的边界"
for i in range(pred.shape[0]):
_place_bybasins_last_forecast(the_array_pred_, i, start, pred[i])
_place_bybasins_last_forecast(the_array_obs_, i, start, obs[i])
preds_xr = valte_dataset.denormalize(the_array_pred_, -1)
obss_xr = valte_dataset.denormalize(the_array_obs_, -1)
else:
# for 3d data, directly process
# TODO: maybe need more test for the pace_idx case
preds_xr = valte_dataset.denormalize(pred)
obss_xr = valte_dataset.denormalize(obs)
def _align_and_order(_obs, _pred):
# 先移除重复的 time 坐标，避免 pandas 索引重复导致对齐失败
def _ensure_unique_time(ds):
if hasattr(ds, "indexes") and "time" in ds.indexes:
idx = ds.indexes["time"]
if getattr(idx, "has_duplicates", False):
keep = ~idx.duplicated()
ds = ds.isel(time=keep)
return ds
_obs = _ensure_unique_time(_obs)
_pred = _ensure_unique_time(_pred)
# 对齐到公共 (basin,time,variable) 的交集，避免 outer 引入 NaN
_obs, _pred = xr.align(_obs, _pred, join="inner")
# time 维为空（无交集）时直接抛错，避免进入 nanmean
if _obs.sizes.get("time", 0) == 0:
raise ValueError(
"No overlapping timestamps between observations and predictions "
f"(obs.time len={_obs.sizes.get('time',0)}, pred.time len={_pred.sizes.get('time',0)})."
)
# 按时间排序（保险）
if "time" in _obs.dims:
_obs = _obs.sortby("time")
if "time" in _pred.dims:
_pred = _pred.sortby("time")
# 规范维度顺序（若存在）
wanted = [d for d in ("basin", "time", "variable") if d in _obs.dims]
_obs = _obs.transpose(*wanted, missing_dims="ignore")
_pred = _pred.transpose(*wanted, missing_dims="ignore")
return _obs, _pred
# 单对象 vs 列表分别处理
if preds_xr is not None and obss_xr is not None:
obss_xr, preds_xr = _align_and_order(obss_xr, preds_xr)
return obss_xr, preds_xr
elif preds_xr_list is not None and obss_xr_list is not None:
obss_aligned, preds_aligned = [], []
for _o, _p in zip(obss_xr_list, preds_xr_list):
_o2, _p2 = _align_and_order(_o, _p)
obss_aligned.append(_o2)
preds_aligned.append(_p2)
return obss_aligned, preds_aligned
else:
# 理论不应走到这
raise RuntimeError("Failed to build preds_xr / obss_xr for evaluation.")
gnn_collate_fn(batch)
¶
Custom collate function for GNN datasets that handles variable-sized graphs.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
batch |
A list of samples, where each sample is a tuple of (sxc, y, edge_index, edge_weight). |
required |
Returns:
| Type | Description |
|---|---|
A list containing batched tensors |
[batched_sxc, batched_y, batched_edge_index, batched_edge_weight, batch_vector] |
Source code in torchhydro/trainers/train_utils.py
def gnn_collate_fn(batch):
"""Custom collate function for GNN datasets that handles variable-sized graphs.
Args:
batch: A list of samples, where each sample is a tuple of
(sxc, y, edge_index, edge_weight).
Returns:
A list containing batched tensors:
[batched_sxc, batched_y, batched_edge_index, batched_edge_weight, batch_vector]
"""
import torch
if len(batch) == 0:
return []
# Unpack the batch
sxc_list, y_list, edge_index_list, edge_weight_list = zip(*batch)
# Batch the target values (y) - these should have the same shape
batched_y = torch.stack(y_list, dim=0) # [batch_size, forecast_length, output_dim]
# Find the maximum number of nodes in this batch
max_num_nodes = max(sxc.shape[0] for sxc in sxc_list)
# Get dimensions
batch_size = len(sxc_list)
seq_length = sxc_list[0].shape[1]
feature_dim = sxc_list[0].shape[2]
# Create padded tensor for node features
batched_sxc = torch.zeros(batch_size, max_num_nodes, seq_length, feature_dim)
# Create batched edge indices and weights
# For each graph in the batch, we need to offset node indices
batched_edge_index = []
batched_edge_weight = []
batch_vector = []
node_offset = 0
for i, (sxc, edge_index, edge_weight) in enumerate(
zip(sxc_list, edge_index_list, edge_weight_list)
):
num_nodes = sxc.shape[0]
# Fill the padded tensor with actual node features
batched_sxc[i, :num_nodes] = sxc
# For edge indices, we need to offset by node_offset to make them unique across batch
if edge_index.numel() > 0:
if edge_index.max() >= num_nodes:
print(
f"Warning: Graph {i} has edge indices {edge_index.max().item()} >= num_nodes {num_nodes}"
)
valid_mask = (edge_index[0] < num_nodes) & (edge_index[1] < num_nodes)
edge_index = edge_index[:, valid_mask]
edge_weight = edge_weight[valid_mask]
if edge_index.numel() > 0:
offset_edge_index = edge_index + node_offset
batched_edge_index.append(offset_edge_index)
batched_edge_weight.append(edge_weight)
# batch_vector: for each node in this graph, assign batch index i
batch_vector.append(torch.full((num_nodes,), i, dtype=torch.long))
node_offset += num_nodes
# Concatenate edge indices and weights if they exist
if batched_edge_index:
batched_edge_index = torch.cat(batched_edge_index, dim=1) # [2, total_edges]
batched_edge_weight = torch.cat(batched_edge_weight, dim=0) # [total_edges]
else:
batched_edge_index = torch.empty((2, 0), dtype=torch.long)
batched_edge_weight = torch.empty(0)
batch_vector = torch.cat(batch_vector, dim=0) # [total_nodes]
return [
batched_sxc,
batched_y,
batched_edge_index,
batched_edge_weight,
batch_vector,
]
mixed_precision_autocast(device, mixed_precision)
¶
Create an autocast context for training/validation.
Source code in torchhydro/trainers/train_utils.py
def mixed_precision_autocast(device: torch.device | None, mixed_precision):
"""Create an autocast context for training/validation."""
dtype = get_mixed_precision_dtype(device, mixed_precision)
if dtype is None:
return nullcontext()
return torch.autocast(device_type=device.type, dtype=dtype)
model_infer(seq_first, device, model, batch, variable_length_cfgs=None, return_key=None, non_blocking=False)
¶
Unified model inference function with variable length support
Parameters¶
seq_first : bool if True, the input data is sequence first device : torch.device cpu or gpu model : torch.nn.Module the model batch : tuple or list batch data from collate_fn or dataset variable_length_cfgs : dict, optional variable length configuration containing mask settings return_key : str, optional when model returns a dict, choose which key (e.g., "f2") to return. if None, defaults to the last frequency (max key).
Source code in torchhydro/trainers/train_utils.py
def model_infer(
seq_first,
device,
model,
batch,
variable_length_cfgs=None,
return_key=None,
non_blocking: bool = False,
):
"""
Unified model inference function with variable length support
Parameters
----------
seq_first : bool
if True, the input data is sequence first
device : torch.device
cpu or gpu
model : torch.nn.Module
the model
batch : tuple or list
batch data from collate_fn or dataset
variable_length_cfgs : dict, optional
variable length configuration containing mask settings
return_key : str, optional
when model returns a dict, choose which key (e.g., "f2") to return.
if None, defaults to the last frequency (max key).
"""
result = get_masked_tensors(variable_length_cfgs, batch, seq_first)
# --- unpack inputs ---
if len(result) == 9:
(
xs,
ys,
edge_index,
edge_weight,
batch_vector,
xs_mask,
ys_mask,
xs_lens,
ys_lens,
) = result
elif len(result) == 8:
xs, ys, edge_index, edge_weight, xs_mask, ys_mask, xs_lens, ys_lens = result
batch_vector = None
else:
xs, ys, xs_mask, ys_mask, xs_lens, ys_lens = result
edge_index = edge_weight = batch_vector = None
# --- move xs to device ---
if isinstance(xs, list):
xs = [
_move_batch_value_to_device(x, device, seq_first, non_blocking)
for x in xs
]
else:
xs = [_move_batch_value_to_device(xs, device, seq_first, non_blocking)]
# --- move ys to device ---
if ys is not None:
ys = (
ys.permute(1, 0, 2).to(device, non_blocking=non_blocking)
if seq_first and ys.ndim == 3
else ys.to(device, non_blocking=non_blocking)
)
# --- move graph data ---
if edge_index is not None:
edge_index = edge_index.to(device, non_blocking=non_blocking)
if edge_weight is not None:
edge_weight = edge_weight.to(device, non_blocking=non_blocking)
if batch_vector is not None:
batch_vector = batch_vector.to(device, non_blocking=non_blocking)
# --- check if model is a generative model (e.g., Diffusion) ---
# Generative models require both input (xs) and target (ys) in forward pass
is_generative = getattr(model, "is_generative", False)
inner_model = getattr(model, "module", model)
# Also check wrapped model (e.g., DistributedDataParallel)
if not is_generative:
is_generative = getattr(inner_model, "is_generative", False)
# --- forward ---
if is_generative:
# Generative models behave differently in train vs eval mode:
# - Train mode: forward(x_condition, y_target) returns {'loss': loss}
# - Eval mode: sample(x_condition) returns generated samples
if model.training:
# Training: compute loss with forward pass
output = inner_model(*xs, ys)
else:
# Inference: generate samples using sample method
output = inner_model.sample(*xs)
elif xs_mask is not None and ys_mask is not None:
if edge_index is not None and edge_weight is not None:
output = model(
*xs,
edge_index=edge_index,
edge_weight=edge_weight,
batch_vector=batch_vector,
mask=xs_mask,
seq_lengths=xs_lens,
)
else:
output = model(*xs, mask=xs_mask, seq_lengths=xs_lens)
else:
if edge_index is not None and edge_weight is not None:
output = model(
*xs,
edge_index=edge_index,
edge_weight=edge_weight,
batch_vector=batch_vector,
)
else:
if edge_index is not None and edge_weight is None:
output = model(
*xs,
edge_index=edge_index,
# batch_vector=batch_vector,
)
else:
output = model(*xs)
# --- handle model outputs ---
if isinstance(output, tuple):
output = output[0]
if isinstance(output, dict):
# Check if this is a generative model that returns loss directly
if "loss" in output:
# Generative models return {'loss': loss_value}, pass it through
output = output
else:
# Multi-frequency output, pick the specified or highest frequency key
if return_key is None:
return_key = sorted(output.keys())[-1] # e.g., "f2"
if return_key not in output:
raise KeyError(
f"Model returned keys {list(output.keys())}, but return_key='{return_key}' not found"
)
output = output[return_key]
if ys_mask is not None:
ys = ys.masked_fill(ys_mask == 0, torch.nan)
# --- seq_first transpose back ---
if seq_first:
# Don't transpose dict outputs (generative models return {'loss': ...} during training)
if isinstance(output, torch.Tensor):
output = output.transpose(0, 1)
if ys is not None:
ys = ys.transpose(0, 1)
return ys, output
torch_single_train(model, opt, criterion, data_loader, device=None, **kwargs)
¶
Training function for one epoch
Parameters¶
model a PyTorch model inherit from nn.Module opt optimizer function from PyTorch optim.Optimizer criterion loss function data_loader object for loading data to the model device where we put the tensors and models
Returns¶
tuple(torch.Tensor, int) loss of this epoch and number of all iterations
Raises¶
ValueError if nan exits, raise a ValueError
Source code in torchhydro/trainers/train_utils.py
def torch_single_train(
model,
opt: optim.Optimizer,
criterion,
data_loader: DataLoader,
device=None,
**kwargs,
):
"""
Training function for one epoch
Parameters
----------
model
a PyTorch model inherit from nn.Module
opt
optimizer function from PyTorch optim.Optimizer
criterion
loss function
data_loader
object for loading data to the model
device
where we put the tensors and models
Returns
-------
tuple(torch.Tensor, int)
loss of this epoch and number of all iterations
Raises
--------
ValueError
if nan exits, raise a ValueError
"""
# we will set model.eval() in the validation function so here we should set model.train()
model.train()
n_iter_ep = 0
which_first_tensor = kwargs["which_first_tensor"]
mixed_precision = kwargs.get("mixed_precision", "off")
seq_first = which_first_tensor != "batch"
variable_length_cfgs = data_loader.dataset.training_cfgs.get(
"variable_length_cfgs", None
)
non_blocking_transfer = bool(
kwargs.get("non_blocking_transfer", kwargs.get("non_blocking", False))
)
torch_device = torch.device(device) if device is not None else torch.device("cpu")
running_loss = 0.0
use_cuda_prefetch = (
bool(kwargs.get("cuda_prefetch", False)) and torch_device.type == "cuda"
)
train_iterable = (
CudaPrefetcher(
data_loader,
torch_device,
non_blocking=non_blocking_transfer,
prefetch_batches=int(kwargs.get("cuda_prefetch_batches", 2)),
performance_monitor=kwargs.get("performance_monitor"),
)
if use_cuda_prefetch
else data_loader
)
pbar = tqdm(train_iterable, total=len(data_loader))
performance_monitor = kwargs.get("performance_monitor")
amp_dtype_name = kwargs.get("amp_dtype")
mixed_mode = _normalize_mixed_precision_mode(mixed_precision)
if amp_dtype_name is None:
amp_dtype_name = "float16" if mixed_mode == "fp16" else "bfloat16"
amp_dtype_name = str(amp_dtype_name).lower()
if amp_dtype_name in {"bfloat16", "bf16"}:
amp_dtype = torch.bfloat16
elif amp_dtype_name in {"float16", "fp16"}:
amp_dtype = torch.float16
else:
raise ValueError("amp_dtype must be bfloat16/bf16 or float16/fp16")
use_amp = (bool(kwargs.get("amp", False)) or mixed_mode != "off") and (
torch_device.type in {"cuda", "cpu"}
)
grad_scaler = torch.amp.GradScaler(
"cuda",
enabled=use_amp and amp_dtype == torch.float16 and torch_device.type == "cuda",
)
nonfinite_check_interval = max(0, int(kwargs.get("nonfinite_check_interval", 100)))
train_iterator = iter(pbar)
while True:
wait_started = time.perf_counter()
try:
batch = next(train_iterator)
except StopIteration:
break
if performance_monitor is not None:
performance_monitor.record_data_wait(time.perf_counter() - wait_started)
compute_events = performance_monitor.cuda_event_pair()
if compute_events is not None:
compute_events[0].record(torch.cuda.current_stream(torch_device))
else:
compute_events = None
compute_started = time.perf_counter()
opt.zero_grad(set_to_none=True)
with torch.autocast(
device_type=torch_device.type,
dtype=amp_dtype,
enabled=use_amp,
):
# mask handling is already done inside model_infer function
trg, output = model_infer(
seq_first,
device,
model,
batch,
variable_length_cfgs,
non_blocking=non_blocking_transfer,
)
batch_vector = getattr(batch, "batch", None)
node_mask = getattr(batch, "node_mask", None)
loss = compute_loss(
trg,
output,
criterion,
batch_vector=batch_vector,
node_mask=node_mask,
**kwargs,
)
# Add auxiliary loss from MoE models (if available)
if hasattr(model, "get_aux_loss"):
aux_loss = model.get_aux_loss()
if aux_loss is not None and aux_loss.device != loss.device:
aux_loss = aux_loss.to(loss.device)
loss = loss + aux_loss
# One scalar read handles validation, warnings, and epoch loss
# accumulation. Calling bool() or item() on several CUDA tensors
# separately introduces repeated host/device synchronization.
loss_value = loss.detach().item()
_validate_training_loss_value(loss_value)
grad_scaler.scale(loss).backward()
grad_scaler.step(opt)
grad_scaler.update()
if compute_events is not None:
compute_events[1].record(torch.cuda.current_stream(torch_device))
performance_monitor.record_compute_events(compute_events)
if performance_monitor is not None:
performance_monitor.record_batch(
batch, time.perf_counter() - compute_started
)
running_loss += loss_value
n_iter_ep += 1
if (
nonfinite_check_interval
and n_iter_ep % nonfinite_check_interval == 0
and not math.isfinite(running_loss)
):
raise ValueError(
"Non-finite training loss detected. Check normalization and input data."
)
if n_iter_ep == 0:
raise ValueError(
"All batch computations of loss result in NAN. Please check the data."
)
if not math.isfinite(running_loss):
raise ValueError(
"Non-finite training loss detected. Check normalization and input data."
)
return running_loss / float(n_iter_ep), n_iter_ep
varied_length_collate_fn(batch)
¶
Collate function for variable length training
This function is automatically used by DataLoader when variable_length_cfgs["use_variable_length"] is True. It pads sequences to the same length and generates corresponding masks.
Parameters¶
batch : list of tuples The batch data after the dataset getitem method
Returns¶
list [xs_pad, ys_pad, xs_lens, ys_lens, xs_mask, ys_mask] - xs_pad: padded input sequences [batch, max_seq_len, input_dim] - ys_pad: padded output sequences [batch, max_seq_len, output_dim] - xs_lens: original sequence lengths for input - ys_lens: original sequence lengths for output - xs_mask: valid position mask for input [batch, max_seq_len] - ys_mask: valid position mask for output [batch, max_seq_len]
Source code in torchhydro/trainers/train_utils.py
def varied_length_collate_fn(batch):
"""Collate function for variable length training
This function is automatically used by DataLoader when variable_length_cfgs["use_variable_length"] is True.
It pads sequences to the same length and generates corresponding masks.
Parameters
----------
batch : list of tuples
The batch data after the dataset __getitem__ method
Returns
-------
list
[xs_pad, ys_pad, xs_lens, ys_lens, xs_mask, ys_mask]
- xs_pad: padded input sequences [batch, max_seq_len, input_dim]
- ys_pad: padded output sequences [batch, max_seq_len, output_dim]
- xs_lens: original sequence lengths for input
- ys_lens: original sequence lengths for output
- xs_mask: valid position mask for input [batch, max_seq_len]
- ys_mask: valid position mask for output [batch, max_seq_len]
"""
xs, ys = zip(*batch)
# sometimes x is a tuple like in dpl dataset, then we can get the shape of the first element as the length
xs_lens = [x[0].shape[0] if type(x) in [tuple, list] else x.shape[0] for x in xs]
ys_lens = [y[0].shape[0] if type(y) in [tuple, list] else y.shape[0] for y in ys]
# if all ys_lens are the same, use default collate_fn to create tensors
if len(set(ys_lens)) == 1 and len(set(xs_lens)) == 1:
xs_tensor = default_collate(xs)
ys_tensor = default_collate(ys)
return [xs_tensor, ys_tensor, None, None, None, None]
# pad the batch data with padding value 0
xs_pad = pad_sequence(xs, batch_first=True, padding_value=0)
ys_pad = pad_sequence(ys, batch_first=True, padding_value=0)
# generate the mask for the batch data
# xs_mask: [batch_size, max_seq_len] or [batch_size, max_seq_len, 1]
batch_size = len(xs_lens)
max_xs_len = max(xs_lens)
max_ys_len = max(ys_lens)
# create the mask for the input sequence (True for valid positions, False for padding positions)
xs_mask = torch.zeros(batch_size, max_xs_len, dtype=torch.bool)
for i, length in enumerate(xs_lens):
xs_mask[i, :length] = True
# create the mask for the output sequence
ys_mask = torch.zeros(batch_size, max_ys_len, dtype=torch.bool)
for i, length in enumerate(ys_lens):
ys_mask[i, :length] = True
return [
xs_pad,
ys_pad,
xs_lens,
ys_lens,
xs_mask,
ys_mask,
]
Resulter
¶
Source code in torchhydro/trainers/resulter.py
class Resulter:
def __init__(self, cfgs) -> None:
self.cfgs = cfgs
self.result_dir = cfgs["data_cfgs"]["case_dir"]
if not os.path.exists(self.result_dir):
os.makedirs(self.result_dir)
@property
def pred_name(self):
return f"epoch{str(self.chosen_trained_epoch)}flow_pred"
@property
def obs_name(self):
return f"epoch{str(self.chosen_trained_epoch)}flow_obs"
@property
def chosen_trained_epoch(self):
model_loader = self.cfgs["evaluation_cfgs"]["model_loader"]
if model_loader["load_way"] == "specified":
epoch_name = str(model_loader["test_epoch"])
elif model_loader["load_way"] == "best":
# NOTE: TO make it consistent with the name in case of model_loader["load_way"] == "pth", the name have to be "best_model.pth"
epoch_name = "best_model.pth"
elif model_loader["load_way"] == "latest":
epoch_name = str(self.cfgs["training_cfgs"]["epochs"])
elif model_loader["load_way"] == "pth":
epoch_name = model_loader["pth_path"].split(os.sep)[-1]
else:
raise ValueError("Invalid load_way")
return epoch_name
def save_cfg(self, cfgs):
# save the cfgs after training
# update the cfgs with the latest one
self.cfgs = cfgs
param_file_exist = any(
(
fnmatch.fnmatch(file, "*.json")
and "_stat" not in file # statistics json file
and "_dict" not in file # data cache json file
)
for file in os.listdir(self.result_dir)
)
if not param_file_exist:
# although we save params log during training, but sometimes we directly evaluate a model
# so here we still save params log if param file does not exist
# no param file was saved yet, here we save data and params setting
save_model_params_log(cfgs, self.result_dir)
def save_result(self, pred, obs):
"""
save the pred value of testing period and obs value
Parameters
----------
pred
predictions
obs
observations
pred_name
the file name of predictions
obs_name
the file name of observations
Returns
-------
None
"""
save_dir = self.result_dir
flow_pred_file = os.path.join(save_dir, self.pred_name)
flow_obs_file = os.path.join(save_dir, self.obs_name)
if isinstance(pred, list) and isinstance(obs, list):
# Case for recover_mode="byforecast", where we have a list of DataArrays per horizon
for i, (p, o) in enumerate(zip(pred, obs)):
max_len = max(len(basin) for basin in p.basin.values)
encoding = {"basin": {"dtype": f"U{max_len}"}}
p.to_netcdf(f"{flow_pred_file}_horizon{i+1}.nc", encoding=encoding)
o.to_netcdf(f"{flow_obs_file}_horizon{i+1}.nc", encoding=encoding)
else:
# Standard case
max_len = max(len(basin) for basin in pred.basin.values)
encoding = {"basin": {"dtype": f"U{max_len}"}}
pred.to_netcdf(flow_pred_file + ".nc", encoding=encoding)
obs.to_netcdf(flow_obs_file + ".nc", encoding=encoding)
def eval_result(self, preds_xr, obss_xr):
# Handle list input for byforecast mode
if isinstance(preds_xr, list) and isinstance(obss_xr, list):
# Evaluate each horizon separately
for i, (p, o) in enumerate(zip(preds_xr, obss_xr)):
# We can reuse the same logic, but maybe append horizon to logs or filenames
# For simplicity, we just evaluate the first horizon or average?
# Usually we want metrics per horizon.
# Here we temporarily adapt to evaluate each one but overwrite or use simple logic
# A better approach: The user might want to see metrics for EACH horizon.
# Let's modify the internal logic to handle this loop if needed,
# OR, for now, we just evaluate them sequentially and print/save separately.
print(f"Evaluating Horizon {i+1}...")
self._eval_single_result(p, o, suffix=f"_horizon{i+1}")
else:
self._eval_single_result(preds_xr, obss_xr)
def _eval_single_result(self, preds_xr, obss_xr, suffix=""):
# types of observations
target_col = self.cfgs["data_cfgs"]["target_cols"]
evaluation_metrics = self.cfgs["evaluation_cfgs"]["metrics"]
basin_ids = self.cfgs["data_cfgs"]["object_ids"]
test_path = self.cfgs["data_cfgs"]["case_dir"]
# Assume object_ids like ['changdian_61561']
# fill_nan: "no" means ignoring the NaN value;
# "sum" means calculate the sum of the following values in the NaN locations.
# For example, observations are [1, nan, nan, 2], and predictions are [0.3, 0.3, 0.3, 1.5].
# Then, "no" means [1, 2] v.s. [0.3, 1.5] while "sum" means [1, 2] v.s. [0.3 + 0.3 + 0.3, 1.5].
# If it is a str, then all target vars use same fill_nan method;
# elif it is a list, each for a var
fill_nan = self.cfgs["evaluation_cfgs"]["fill_nan"]
# Then evaluate the model metrics
if isinstance(fill_nan, list) and len(fill_nan) != len(target_col):
raise ValueError("length of fill_nan must be equal to target_col's")
for i, col in enumerate(target_col):
eval_log = {}
obs = obss_xr[col].to_numpy()
pred = preds_xr[col].to_numpy()
eval_log = calculate_and_record_metrics(
obs,
pred,
evaluation_metrics,
col,
fill_nan[i] if isinstance(fill_nan, list) else fill_nan,
eval_log,
)
# Create pandas DataFrames from eval_log for each target variable (e.g., streamflow)
# Create a dictionary to hold the data for the DataFrame
data = {}
# Iterate over metrics in eval_log
for metric, values in eval_log.items():
# Remove 'of streamflow' (or similar) from the metric name
clean_metric = metric.replace(f"of {col}", "").strip()
# Add the cleaned metric to the data dictionary
data[clean_metric] = values
# Create a DataFrame using object_ids as the index and metrics as columns
df = pd.DataFrame(data, index=basin_ids)
# Save the DataFrame to a CSV file
output_file = os.path.join(test_path, f"metric_{col}{suffix}.csv")
df.to_csv(output_file, index_label="basin_id")
# Finally, try to explain model behaviour using shap
is_shap = self.cfgs["evaluation_cfgs"]["explainer"] == "shap"
if is_shap and suffix == "": # Only run SHAP once or for main result
shap_summary_plot(self.model, self.traindataset, self.testdataset)
# deep_explain_model_summary_plot(self.model, test_data)
# deep_explain_model_heatmap(self.model, test_data)
def _convert_streamflow_units(self, ds):
"""convert the streamflow units to m^3/s
Parameters
----------
ds : xr.Dataset
xarray Dataset containing predictions or observations
Returns
-------
xr.Dataset
"""
data_cfgs = self.cfgs["data_cfgs"]
from torchhydro.configs.data_resolver import open_dataset_from_source_cfgs
sc = data_cfgs["source_cfgs"]
data_source = open_dataset_from_source_cfgs(sc)
basin_id = data_cfgs["object_ids"]
try:
basin_area = data_source.read_area(basin_id)
except (AttributeError, NotImplementedError) as e:
raise AttributeError(
f"Data source '{sc.get('dataset_id')}' does not implement "
f"read_area(), which is required for streamflow unit conversion."
) from e
target_unit = "m^3/s"
target_cols = data_cfgs["target_cols"]
if not target_cols:
raise ValueError(
"target_cols is empty — cannot determine the flow variable "
"for unit conversion. Add a flow variable (e.g. 'streamflow') "
"to target_cols."
)
var_flow = target_cols[0]
streamflow_ds = ds[[var_flow]]
ds_ = streamflow_unit_conv(
streamflow_ds, basin_area, target_unit=target_unit, inverse=True
)
new_ds = ds.copy(deep=True)
new_ds[var_flow] = ds_[var_flow]
return new_ds
def load_result(self, convert_flow_unit=False) -> Tuple[np.array, np.array]:
"""load the pred value of testing period and obs value"""
save_dir = self.result_dir
pred_file = os.path.join(save_dir, self.pred_name + ".nc")
obs_file = os.path.join(save_dir, self.obs_name + ".nc")
pred = xr.open_dataset(pred_file)
obs = xr.open_dataset(obs_file)
if convert_flow_unit:
pred = self._convert_streamflow_units(pred)
obs = self._convert_streamflow_units(obs)
return pred, obs
def save_intermediate_results(self, **kwargs):
"""Load model weights and deal with some intermediate results"""
is_cell_states = kwargs.get("is_cell_states", False)
is_pbm_params = kwargs.get("is_pbm_params", False)
cfgs = self.cfgs
cfgs["training_cfgs"]["train_mode"] = False
training_cfgs = cfgs["training_cfgs"]
seq_first = training_cfgs["which_first_tensor"] == "sequence"
if is_cell_states:
raise NotImplementedError(
"return_cell_states is not supported yet "
"(save_intermediate_results)"
)
if is_pbm_params:
self._save_pbm_params(cfgs, seq_first)
def _save_pbm_params(self, cfgs, seq_first):
training_cfgs = cfgs["training_cfgs"]
model_loader = cfgs["evaluation_cfgs"]["model_loader"]
model_pth_dir = cfgs["data_cfgs"]["case_dir"]
weight_path = read_pth_from_model_loader(model_loader, model_pth_dir)
cfgs["model_cfgs"]["weight_path"] = weight_path
cfgs["training_cfgs"]["device"] = [0] if torch.cuda.is_available() else [-1]
deephydro = DeepHydro(cfgs)
device = deephydro.device
dl_model = deephydro.model.dl_model
pb_model = deephydro.model.pb_model
param_func = deephydro.model.param_func
# TODO: check for dplnnmodule model
param_test_way = deephydro.model.param_test_way
test_dataloader = DataLoader(
deephydro.testdataset,
batch_size=training_cfgs["batch_size"],
shuffle=False,
sampler=None,
batch_sampler=None,
drop_last=False,
timeout=0,
worker_init_fn=None,
)
deephydro.model.eval()
# here the batch is just an index of lookup table, so any batch size could be chosen
params_lst = []
with torch.no_grad():
for batch in test_dataloader:
ys, gen = model_infer(seq_first, device, dl_model, batch)
# we set all params' values in [0, 1] and will scale them when forwarding
if param_func == "clamp":
params_ = torch.clamp(gen, min=0.0, max=1.0)
elif param_func == "sigmoid":
params_ = F.sigmoid(gen)
else:
raise NotImplementedError(
"We don't provide this way to limit parameters' range!! Please choose sigmoid or clamp"
)
# just get one-period values, here we use the final period's values
params = params_[:, -1, :]
params_lst.append(params)
pb_params = reduce(lambda a, b: torch.cat((a, b), dim=0), params_lst)
# trans tensor to pandas dataframe
sites = deephydro.cfgs["data_cfgs"]["object_ids"]
params_names = pb_model.params_names
params_df = pd.DataFrame(
pb_params.cpu().numpy(), columns=params_names, index=sites
)
save_param_file = os.path.join(
model_pth_dir, f"pb_params_{int(time.time())}.csv"
)
params_df.to_csv(save_param_file, index_label="GAGE_ID")
def read_tensorboard_log(self, **kwargs):
"""read tensorboard log files"""
is_scalar = kwargs.get("is_scalar", False)
is_histogram = kwargs.get("is_histogram", False)
log_dir = self.cfgs["data_cfgs"]["case_dir"]
if is_scalar:
scalar_file = os.path.join(log_dir, "tb_scalars.csv")
if not os.path.exists(scalar_file):
reader = SummaryReader(log_dir)
df_scalar = reader.scalars
df_scalar.to_csv(scalar_file, index=False)
else:
df_scalar = pd.read_csv(scalar_file)
if is_histogram:
histogram_file = os.path.join(log_dir, "tb_histograms.csv")
if not os.path.exists(histogram_file):
reader = SummaryReader(log_dir)
df_histogram = reader.histograms
df_histogram.to_csv(histogram_file, index=False)
else:
df_histogram = pd.read_csv(histogram_file)
if is_scalar and is_histogram:
return df_scalar, df_histogram
elif is_scalar:
return df_scalar
elif is_histogram:
return df_histogram
# TODO: the following code is not finished yet
def load_ensemble_result(
self, save_dirs, test_epoch, flow_unit="m3/s", basin_areas=None
) -> Tuple[np.array, np.array]:
"""
load ensemble mean value
Parameters
----------
save_dirs
test_epoch
flow_unit
default is m3/s, if it is not m3/s, transform the results
basin_areas
if unit is mm/day it will be used, default is None
Returns
-------
"""
preds = []
obss = []
for save_dir in save_dirs:
pred_i, obs_i = self.load_result(save_dir, test_epoch)
if pred_i.ndim == 3 and pred_i.shape[-1] == 1:
pred_i = pred_i.reshape(pred_i.shape[0], pred_i.shape[1])
obs_i = obs_i.reshape(obs_i.shape[0], obs_i.shape[1])
preds.append(pred_i)
obss.append(obs_i)
preds_np = np.array(preds)
obss_np = np.array(obss)
pred_mean = np.mean(preds_np, axis=0)
obs_mean = np.mean(obss_np, axis=0)
if flow_unit == "mm/day":
if basin_areas is None:
raise ArithmeticError("No basin areas we cannot calculate")
basin_areas = np.repeat(basin_areas, obs_mean.shape[1], axis=0).reshape(
obs_mean.shape
)
obs_mean = obs_mean * basin_areas * 1e-3 * 1e6 / 86400
pred_mean = pred_mean * basin_areas * 1e-3 * 1e6 / 86400
elif flow_unit == "m3/s":
pass
elif flow_unit == "ft3/s":
obs_mean = obs_mean / 35.314666721489
pred_mean = pred_mean / 35.314666721489
return pred_mean, obs_mean
def eval_ensemble_result(
self,
save_dirs,
test_epoch,
return_value=False,
flow_unit="m3/s",
basin_areas=None,
) -> Tuple[np.array, np.array]:
"""calculate statistics for ensemble results
Parameters
----------
save_dirs : _type_
where the results save
test_epoch : _type_
we name the results files with the test_epoch
return_value : bool, optional
if True, return (inds_df, pred_mean, obs_mean), by default False
flow_unit : str, optional
arg for load_ensemble_result, by default "m3/s"
basin_areas : _type_, optional
arg for load_ensemble_result, by default None
Returns
-------
Tuple[np.array, np.array]
inds_df or (inds_df, pred_mean, obs_mean)
"""
pred_mean, obs_mean = self.load_ensemble_result(
save_dirs, test_epoch, flow_unit=flow_unit, basin_areas=basin_areas
)
inds = stat_error(obs_mean, pred_mean)
inds_df = pd.DataFrame(inds)
return (inds_df, pred_mean, obs_mean) if return_value else inds_df
eval_ensemble_result(self, save_dirs, test_epoch, return_value=False, flow_unit='m3/s', basin_areas=None)
¶
calculate statistics for ensemble results
Parameters¶
save_dirs : type where the results save test_epoch : type we name the results files with the test_epoch return_value : bool, optional if True, return (inds_df, pred_mean, obs_mean), by default False flow_unit : str, optional arg for load_ensemble_result, by default "m3/s" basin_areas : type, optional arg for load_ensemble_result, by default None
Returns¶
Tuple[np.array, np.array] inds_df or (inds_df, pred_mean, obs_mean)
Source code in torchhydro/trainers/resulter.py
def eval_ensemble_result(
self,
save_dirs,
test_epoch,
return_value=False,
flow_unit="m3/s",
basin_areas=None,
) -> Tuple[np.array, np.array]:
"""calculate statistics for ensemble results
Parameters
----------
save_dirs : _type_
where the results save
test_epoch : _type_
we name the results files with the test_epoch
return_value : bool, optional
if True, return (inds_df, pred_mean, obs_mean), by default False
flow_unit : str, optional
arg for load_ensemble_result, by default "m3/s"
basin_areas : _type_, optional
arg for load_ensemble_result, by default None
Returns
-------
Tuple[np.array, np.array]
inds_df or (inds_df, pred_mean, obs_mean)
"""
pred_mean, obs_mean = self.load_ensemble_result(
save_dirs, test_epoch, flow_unit=flow_unit, basin_areas=basin_areas
)
inds = stat_error(obs_mean, pred_mean)
inds_df = pd.DataFrame(inds)
return (inds_df, pred_mean, obs_mean) if return_value else inds_df
load_ensemble_result(self, save_dirs, test_epoch, flow_unit='m3/s', basin_areas=None)
¶
load ensemble mean value
Parameters¶
save_dirs test_epoch flow_unit default is m3/s, if it is not m3/s, transform the results basin_areas if unit is mm/day it will be used, default is None
Returns¶
Source code in torchhydro/trainers/resulter.py
def load_ensemble_result(
self, save_dirs, test_epoch, flow_unit="m3/s", basin_areas=None
) -> Tuple[np.array, np.array]:
"""
load ensemble mean value
Parameters
----------
save_dirs
test_epoch
flow_unit
default is m3/s, if it is not m3/s, transform the results
basin_areas
if unit is mm/day it will be used, default is None
Returns
-------
"""
preds = []
obss = []
for save_dir in save_dirs:
pred_i, obs_i = self.load_result(save_dir, test_epoch)
if pred_i.ndim == 3 and pred_i.shape[-1] == 1:
pred_i = pred_i.reshape(pred_i.shape[0], pred_i.shape[1])
obs_i = obs_i.reshape(obs_i.shape[0], obs_i.shape[1])
preds.append(pred_i)
obss.append(obs_i)
preds_np = np.array(preds)
obss_np = np.array(obss)
pred_mean = np.mean(preds_np, axis=0)
obs_mean = np.mean(obss_np, axis=0)
if flow_unit == "mm/day":
if basin_areas is None:
raise ArithmeticError("No basin areas we cannot calculate")
basin_areas = np.repeat(basin_areas, obs_mean.shape[1], axis=0).reshape(
obs_mean.shape
)
obs_mean = obs_mean * basin_areas * 1e-3 * 1e6 / 86400
pred_mean = pred_mean * basin_areas * 1e-3 * 1e6 / 86400
elif flow_unit == "m3/s":
pass
elif flow_unit == "ft3/s":
obs_mean = obs_mean / 35.314666721489
pred_mean = pred_mean / 35.314666721489
return pred_mean, obs_mean
load_result(self, convert_flow_unit=False)
¶
load the pred value of testing period and obs value
Source code in torchhydro/trainers/resulter.py
def load_result(self, convert_flow_unit=False) -> Tuple[np.array, np.array]:
"""load the pred value of testing period and obs value"""
save_dir = self.result_dir
pred_file = os.path.join(save_dir, self.pred_name + ".nc")
obs_file = os.path.join(save_dir, self.obs_name + ".nc")
pred = xr.open_dataset(pred_file)
obs = xr.open_dataset(obs_file)
if convert_flow_unit:
pred = self._convert_streamflow_units(pred)
obs = self._convert_streamflow_units(obs)
return pred, obs
read_tensorboard_log(self, **kwargs)
¶
read tensorboard log files
Source code in torchhydro/trainers/resulter.py
def read_tensorboard_log(self, **kwargs):
"""read tensorboard log files"""
is_scalar = kwargs.get("is_scalar", False)
is_histogram = kwargs.get("is_histogram", False)
log_dir = self.cfgs["data_cfgs"]["case_dir"]
if is_scalar:
scalar_file = os.path.join(log_dir, "tb_scalars.csv")
if not os.path.exists(scalar_file):
reader = SummaryReader(log_dir)
df_scalar = reader.scalars
df_scalar.to_csv(scalar_file, index=False)
else:
df_scalar = pd.read_csv(scalar_file)
if is_histogram:
histogram_file = os.path.join(log_dir, "tb_histograms.csv")
if not os.path.exists(histogram_file):
reader = SummaryReader(log_dir)
df_histogram = reader.histograms
df_histogram.to_csv(histogram_file, index=False)
else:
df_histogram = pd.read_csv(histogram_file)
if is_scalar and is_histogram:
return df_scalar, df_histogram
elif is_scalar:
return df_scalar
elif is_histogram:
return df_histogram
save_intermediate_results(self, **kwargs)
¶
Load model weights and deal with some intermediate results
Source code in torchhydro/trainers/resulter.py
def save_intermediate_results(self, **kwargs):
"""Load model weights and deal with some intermediate results"""
is_cell_states = kwargs.get("is_cell_states", False)
is_pbm_params = kwargs.get("is_pbm_params", False)
cfgs = self.cfgs
cfgs["training_cfgs"]["train_mode"] = False
training_cfgs = cfgs["training_cfgs"]
seq_first = training_cfgs["which_first_tensor"] == "sequence"
if is_cell_states:
raise NotImplementedError(
"return_cell_states is not supported yet "
"(save_intermediate_results)"
)
if is_pbm_params:
self._save_pbm_params(cfgs, seq_first)
save_result(self, pred, obs)
¶
save the pred value of testing period and obs value
Parameters¶
pred predictions obs observations pred_name the file name of predictions obs_name the file name of observations
Returns¶
None
Source code in torchhydro/trainers/resulter.py
def save_result(self, pred, obs):
"""
save the pred value of testing period and obs value
Parameters
----------
pred
predictions
obs
observations
pred_name
the file name of predictions
obs_name
the file name of observations
Returns
-------
None
"""
save_dir = self.result_dir
flow_pred_file = os.path.join(save_dir, self.pred_name)
flow_obs_file = os.path.join(save_dir, self.obs_name)
if isinstance(pred, list) and isinstance(obs, list):
# Case for recover_mode="byforecast", where we have a list of DataArrays per horizon
for i, (p, o) in enumerate(zip(pred, obs)):
max_len = max(len(basin) for basin in p.basin.values)
encoding = {"basin": {"dtype": f"U{max_len}"}}
p.to_netcdf(f"{flow_pred_file}_horizon{i+1}.nc", encoding=encoding)
o.to_netcdf(f"{flow_obs_file}_horizon{i+1}.nc", encoding=encoding)
else:
# Standard case
max_len = max(len(basin) for basin in pred.basin.values)
encoding = {"basin": {"dtype": f"U{max_len}"}}
pred.to_netcdf(flow_pred_file + ".nc", encoding=encoding)
obs.to_netcdf(flow_obs_file + ".nc", encoding=encoding)
Author: Wenyu Ouyang Date: 2023-07-25 16:47:19 LastEditTime: 2025-06-17 10:39:32 LastEditors: Wenyu Ouyang Description: Lightning Fabric wrapper for debugging and distributed training FilePath: orchhydro orchhydro rainers abric_wrapper.py Copyright (c) 2025-2026 Wenyu Ouyang. All rights reserved.
FabricWrapper
¶
A wrapper class that can switch between Lightning Fabric and normal PyTorch operations based on configuration settings.
TODO: the fabric wrapper is not fully used for parallel training yet
Source code in torchhydro/trainers/fabric_wrapper.py
class FabricWrapper:
"""
A wrapper class that can switch between Lightning Fabric and normal PyTorch operations
based on configuration settings.
TODO: the fabric wrapper is not fully used for parallel training yet
"""
def __init__(self, use_fabric: bool = True, fabric_config: Optional[Dict] = None):
"""
Initialize the Fabric wrapper.
Parameters
----------
use_fabric : bool
Whether to use Lightning Fabric or normal PyTorch operations
fabric_config : Optional[Dict]
Configuration for Fabric (devices, strategy, etc.)
"""
self.use_fabric = use_fabric
self.fabric_config = fabric_config or {}
self._fabric: Optional[Any] = None
self._device: Optional[torch.device] = None
if self.use_fabric:
self._init_fabric()
else:
self._init_pytorch()
def _init_fabric(self) -> None:
"""Initialize Lightning Fabric"""
try:
import lightning as L
# Default fabric configuration
default_config = {
"accelerator": "auto",
"devices": "auto",
"strategy": "auto",
"precision": "32-true",
}
# Update with user config
default_config.update(self.fabric_config)
self._fabric = L.Fabric(**default_config)
print("✅ Lightning Fabric initialized successfully")
except ImportError:
print("❌ Lightning not found, falling back to normal PyTorch")
self.use_fabric = False
self._init_pytorch()
def _init_pytorch(self) -> None:
"""Initialize normal PyTorch setup"""
self.device_num = self.fabric_config["devices"]
# self.device_num = [0]
self._device = get_the_device(self.device_num)
print(f"[OK] Normal PyTorch initialized, using device: {self._device}")
def setup_module(self, model: torch.nn.Module) -> torch.nn.Module:
"""Setup model for training"""
if self.use_fabric:
return self._fabric.setup_module(model)
else:
return model.to(self._device)
def setup_optimizers(
self, optimizer: torch.optim.Optimizer
) -> torch.optim.Optimizer:
"""Setup optimizer"""
if self.use_fabric:
return self._fabric.setup_optimizers(optimizer)
else:
return optimizer
def setup_dataloaders(
self, *dataloaders: torch.utils.data.DataLoader
) -> Tuple[torch.utils.data.DataLoader, ...]:
"""Setup dataloaders"""
if self.use_fabric:
return self._fabric.setup_dataloaders(*dataloaders)
else:
return dataloaders
def save(self, path: str, state_dict: Dict[str, Any]) -> None:
"""Save model state"""
if self.use_fabric:
self._fabric.save(path, state_dict)
else:
torch.save(state_dict, path)
def load(self, path: str, model: Optional[torch.nn.Module] = None) -> Any:
"""Load model state"""
if self.use_fabric:
return self._fabric.load(path, model)
else:
return torch.load(path, map_location=self._device)
def load_raw(self, path: str, model: torch.nn.Module) -> None:
"""Load raw model weights"""
if self.use_fabric:
checkpoint = self._fabric.load(path)
model.load_state_dict(checkpoint)
else:
checkpoint = torch.load(path, map_location=self._device)
model.load_state_dict(checkpoint)
def launch(self, fn: Optional[Any] = None, *args: Any, **kwargs: Any) -> Any:
"""Launch training function"""
if self.use_fabric:
if fn is None:
# This is called without a function, just launch fabric
return self._fabric.launch()
else:
return self._fabric.launch(fn, *args, **kwargs)
else:
# Normal PyTorch, just call the function directly
if fn is not None:
return fn(*args, **kwargs)
else:
return None
def backward(self, loss: torch.Tensor) -> None:
"""Backward pass"""
if self.use_fabric:
self._fabric.backward(loss)
else:
loss.backward()
def clip_gradients(
self,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
max_norm: float = 1.0,
) -> None:
"""Clip gradients"""
if self.use_fabric:
self._fabric.clip_gradients(model, optimizer, max_norm=max_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
@property
def device(self) -> torch.device:
"""Get current device"""
if self.use_fabric:
return self._fabric.device
else:
return self._device
@property
def local_rank(self) -> int:
"""Get local rank"""
if self.use_fabric:
return self._fabric.local_rank
else:
return 0
@property
def global_rank(self) -> int:
"""Get global rank"""
if self.use_fabric:
return self._fabric.global_rank
else:
return 0
@property
def world_size(self) -> int:
"""Get world size"""
if self.use_fabric:
return self._fabric.world_size
else:
return 1
def barrier(self) -> None:
"""Synchronization barrier"""
if self.use_fabric:
self._fabric.barrier()
else:
pass # No barrier needed for single process
def print(self, *args: Any, **kwargs: Any) -> None:
"""Print only on rank 0"""
if self.use_fabric:
self._fabric.print(*args, **kwargs)
else:
print(*args, **kwargs)
device: device
property
readonly
¶
Get current device
global_rank: int
property
readonly
¶
Get global rank
local_rank: int
property
readonly
¶
Get local rank
world_size: int
property
readonly
¶
Get world size
__init__(self, use_fabric=True, fabric_config=None)
special
¶
Initialize the Fabric wrapper.
Parameters¶
use_fabric : bool Whether to use Lightning Fabric or normal PyTorch operations fabric_config : Optional[Dict] Configuration for Fabric (devices, strategy, etc.)
Source code in torchhydro/trainers/fabric_wrapper.py
def __init__(self, use_fabric: bool = True, fabric_config: Optional[Dict] = None):
"""
Initialize the Fabric wrapper.
Parameters
----------
use_fabric : bool
Whether to use Lightning Fabric or normal PyTorch operations
fabric_config : Optional[Dict]
Configuration for Fabric (devices, strategy, etc.)
"""
self.use_fabric = use_fabric
self.fabric_config = fabric_config or {}
self._fabric: Optional[Any] = None
self._device: Optional[torch.device] = None
if self.use_fabric:
self._init_fabric()
else:
self._init_pytorch()
backward(self, loss)
¶
Backward pass
Source code in torchhydro/trainers/fabric_wrapper.py
def backward(self, loss: torch.Tensor) -> None:
"""Backward pass"""
if self.use_fabric:
self._fabric.backward(loss)
else:
loss.backward()
barrier(self)
¶
Synchronization barrier
Source code in torchhydro/trainers/fabric_wrapper.py
def barrier(self) -> None:
"""Synchronization barrier"""
if self.use_fabric:
self._fabric.barrier()
else:
pass # No barrier needed for single process
clip_gradients(self, model, optimizer, max_norm=1.0)
¶
Clip gradients
Source code in torchhydro/trainers/fabric_wrapper.py
def clip_gradients(
self,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
max_norm: float = 1.0,
) -> None:
"""Clip gradients"""
if self.use_fabric:
self._fabric.clip_gradients(model, optimizer, max_norm=max_norm)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
launch(self, fn=None, *args, **kwargs)
¶
Launch training function
Source code in torchhydro/trainers/fabric_wrapper.py
def launch(self, fn: Optional[Any] = None, *args: Any, **kwargs: Any) -> Any:
"""Launch training function"""
if self.use_fabric:
if fn is None:
# This is called without a function, just launch fabric
return self._fabric.launch()
else:
return self._fabric.launch(fn, *args, **kwargs)
else:
# Normal PyTorch, just call the function directly
if fn is not None:
return fn(*args, **kwargs)
else:
return None
load(self, path, model=None)
¶
Load model state
Source code in torchhydro/trainers/fabric_wrapper.py
def load(self, path: str, model: Optional[torch.nn.Module] = None) -> Any:
"""Load model state"""
if self.use_fabric:
return self._fabric.load(path, model)
else:
return torch.load(path, map_location=self._device)
load_raw(self, path, model)
¶
Load raw model weights
Source code in torchhydro/trainers/fabric_wrapper.py
def load_raw(self, path: str, model: torch.nn.Module) -> None:
"""Load raw model weights"""
if self.use_fabric:
checkpoint = self._fabric.load(path)
model.load_state_dict(checkpoint)
else:
checkpoint = torch.load(path, map_location=self._device)
model.load_state_dict(checkpoint)
print(self, *args, **kwargs)
¶
Print only on rank 0
Source code in torchhydro/trainers/fabric_wrapper.py
def print(self, *args: Any, **kwargs: Any) -> None:
"""Print only on rank 0"""
if self.use_fabric:
self._fabric.print(*args, **kwargs)
else:
print(*args, **kwargs)
save(self, path, state_dict)
¶
Save model state
Source code in torchhydro/trainers/fabric_wrapper.py
def save(self, path: str, state_dict: Dict[str, Any]) -> None:
"""Save model state"""
if self.use_fabric:
self._fabric.save(path, state_dict)
else:
torch.save(state_dict, path)
setup_dataloaders(self, *dataloaders)
¶
Setup dataloaders
Source code in torchhydro/trainers/fabric_wrapper.py
def setup_dataloaders(
self, *dataloaders: torch.utils.data.DataLoader
) -> Tuple[torch.utils.data.DataLoader, ...]:
"""Setup dataloaders"""
if self.use_fabric:
return self._fabric.setup_dataloaders(*dataloaders)
else:
return dataloaders
setup_module(self, model)
¶
Setup model for training
Source code in torchhydro/trainers/fabric_wrapper.py
def setup_module(self, model: torch.nn.Module) -> torch.nn.Module:
"""Setup model for training"""
if self.use_fabric:
return self._fabric.setup_module(model)
else:
return model.to(self._device)
setup_optimizers(self, optimizer)
¶
Setup optimizer
Source code in torchhydro/trainers/fabric_wrapper.py
def setup_optimizers(
self, optimizer: torch.optim.Optimizer
) -> torch.optim.Optimizer:
"""Setup optimizer"""
if self.use_fabric:
return self._fabric.setup_optimizers(optimizer)
else:
return optimizer
create_fabric_wrapper(training_cfgs)
¶
Create a fabric wrapper based on training configuration.
Parameters¶
training_cfgs : Dict Training configuration dictionary
Returns¶
FabricWrapper Initialized fabric wrapper
Source code in torchhydro/trainers/fabric_wrapper.py
def create_fabric_wrapper(training_cfgs: Dict) -> FabricWrapper:
"""
Create a fabric wrapper based on training configuration.
Parameters
----------
training_cfgs : Dict
Training configuration dictionary
Returns
-------
FabricWrapper
Initialized fabric wrapper
"""
# Check if we should use fabric
fabric_strategy = training_cfgs.get("fabric_strategy")
use_fabric = fabric_strategy is not None
# Check if we have multiple devices
devices = training_cfgs.get("device", [0])
if isinstance(devices, list) and len(devices) == 1 and use_fabric:
print("📱 Single device detected - we can disable Fabric")
use_fabric = False
# Fabric configuration
if devices == "auto" or devices == ["auto"]:
final_devices = "auto"
else:
final_devices = devices if isinstance(devices, list) else [devices]
fabric_config = {
"devices": final_devices,
"strategy": fabric_strategy,
"precision": training_cfgs.get("precision", "32-true"),
"accelerator": training_cfgs.get("accelerator", "auto"),
}
return FabricWrapper(use_fabric=use_fabric, fabric_config=fabric_config)