PyTorchTrainingPipeline Module

Utilities and pipeline helpers for training PyTorch models (training loops, schedulers, metrics helpers and common pipelines).

HMB.PyTorchTrainingPipeline.TrainEvaluateClassificationModel(model, criterion, device, bestModelStoragePath, noOfClasses, numEpochs, optimizer, scaler=None, scheduler=None, trainLoader=None, valLoader=None, resumeFromCheckpoint=False, finalModelStoragePath=None, judgeBy='both', earlyStopping=None, earlyStoppingPatience=None, checkpointSaver=None, verbose=True, gradAccumSteps=1, maxGradNorm=None, useAmp=True, useMixupFn=False, mixUpAlpha=0.5, useEma=False, saveEvery=None)[source]

Train and evaluate a classification model for a specified number of epochs.

Parameters:

  • model (torch.nn.Module) – Model to train and evaluate.

  • criterion (callable) – Loss function.

  • device (torch.device) – Device to run training and evaluation on (CPU or GPU).

  • bestModelStoragePath (str) – Path to save the best model.

  • noOfClasses (int) – Number of classes in the classification task.

  • numEpochs (int) – Total number of epochs for training.

  • optimizer (torch.optim.Optimizer) – Optimizer for updating model parameters.

  • scaler (torch.amp.GradScaler) – Gradient scaler for mixed precision training.

  • scheduler (torch.optim.lr_scheduler._LRScheduler or torch.optim.lr_scheduler.ReduceLROnPlateau) – Learning rate scheduler.

  • trainLoader (torch.utils.data.DataLoader) – DataLoader for training data.

  • valLoader (torch.utils.data.DataLoader) – DataLoader for validation data.

  • resumeFromCheckpoint (bool, optional) – Flag to indicate if training should resume from a checkpoint. Defaults to False.

  • finalModelStoragePath (str, optional) – Path to save the final model after training. Defaults to None.

  • judgeBy (str, optional) – Criterion to judge the best model (“val_loss”, “val_accuracy”, or “both”). Defaults to “both”.

  • earlyStopping (object, optional) – EarlyStopping callback instance. If None, uses earlyStoppingPatience fallback. Defaults to None.

  • earlyStoppingPatience (int, optional) – Patience for early stopping. Defaults to None.

  • checkpointSaver (object, optional) – CheckpointSaver callback instance for saving best/latest checkpoints. Defaults to None.

  • verbose (bool, optional) – Verbosity flag to control logging. Defaults to True.

  • gradAccumSteps (int, optional) – Number of gradient accumulation steps. Defaults to 1. When >1, gradients are accumulated over multiple batches before performing an optimizer step. When using gradient accumulation, ensure that the effective batch size (batchSize * gradAccumSteps) fits in memory.

  • maxGradNorm (float, optional) – Maximum gradient norm for clipping. Defaults to None.

  • useAmp (bool, optional) – Whether to use automatic mixed precision. Defaults to True.

  • useMixupFn (bool, optional) – Whether to use MixUp data augmentation. Defaults to False.

  • mixUpAlpha (float, optional) – Value of alpha for MixUp data augmentation. Defaults to 0.5.

  • useEma (object, optional) – Whether to use Exponential Moving Average for model parameters. Defaults to False.

  • saveEvery (int, optional) – Save model every N epochs. Defaults to None.

Returns:

History dictionary containing training and validation metrics.

Return type:

dict

Examples

import torch
from torch import nn, optim
from torch.amp import GradScaler
from torch.utils.data import DataLoader
from HMB.PyTorchHelper import TrainEvaluateModel

# Prepare the data loaders.
# Replace with actual dataset and DataLoader code.
trainLoader = DataLoader(...)
valLoader = DataLoader(...)

# Create the model.
# Replace with actual model creation code.
model = pth.CreateTimmModel("resnet18", numClasses=10, pretrained=True)

# Define loss function, optimizer, scaler, and scheduler.
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
scaler = GradScaler()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

# Define device.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Train and evaluate the model.
history = TrainEvaluateModel(
  model,
  criterion,
  device,
  bestModelStoragePath="best_model.pth",
  noOfClasses=10,
  numEpochs=25,
  optimizer=optimizer,
  scaler=scaler,
  scheduler=scheduler,
  trainLoader=trainLoader,
  valLoader=valLoader,
  resumeFromCheckpoint=False,
  finalModelStoragePath="final_model.pth",
  judgeBy="both",
  earlyStoppingPatience=None,
  verbose=True,
  gradAccumSteps=1,
  maxGradNorm=None,
  useAmp=True,
  useMixupFn=False,
  mixUpAlpha=0.5,
  useEma=False,
  saveEvery=None
)
HMB.PyTorchTrainingPipeline.TrainOneEpoch(model, dataLoader, criterion, device, epoch, noOfClasses, numEpochs, optimizer, scaler, gradAccumSteps=1, maxGradNorm=None, useAmp=True, useMixupFn=False, mixUpAlpha=0.5, ema=None, verbose=True)[source]

Train the model for one epoch.

Parameters:

  • model (torch.nn.Module) – Model to train.

  • dataLoader (torch.utils.data.DataLoader) – DataLoader for training data.

  • criterion (callable) – Loss function.

  • device (torch.device) – Device to run training on (CPU or GPU).

  • epoch (int) – Current epoch number.

  • noOfClasses (int) – Number of classes in the classification task.

  • numEpochs (int) – Total number of epochs for training.

  • optimizer (torch.optim.Optimizer) – Optimizer for updating model parameters.

  • scaler (torch.amp.GradScaler) – Gradient scaler for mixed precision training.

  • gradAccumSteps (int, optional) – Number of gradient accumulation steps. Defaults to 1.

  • maxGradNorm (float, optional) – Maximum gradient norm for clipping. Defaults to None.

  • useAmp (bool, optional) – Whether to use automatic mixed precision. Defaults to True.

  • useMixupFn (bool, optional) – Whether to use MixUp data augmentation. Defaults to False.

  • mixUpAlpha (float, optional) – Value of alpha for MixUp data augmentation. Defaults to 0.5.

  • ema (object, optional) – Exponential Moving Average object. Defaults to None.

  • verbose (bool, optional) – Verbosity flag to control logging. Defaults to True.

Returns:

(avgTrainLoss, avgTrainAccuracy) for the epoch.

  • avgTrainLoss (float): Average training loss for the epoch.

  • avgTrainAccuracy (float): Average training accuracy for the epoch.

Return type:

tuple

HMB.PyTorchTrainingPipeline.EvaluateOneEpoch(model, dataLoader, criterion, device, noOfClasses)[source]

Evaluate the model for one epoch.

Parameters:

  • model (torch.nn.Module) – Model to evaluate.

  • dataLoader (torch.utils.data.DataLoader) – DataLoader for evaluation data.

  • criterion (callable) – Loss function.

  • device (torch.device) – Device to run evaluation on (CPU or GPU).

  • noOfClasses (int) – Number of classes in the classification task.

Returns:

(avgValLoss, avgValAccuracy) for the epoch.

  • avgValLoss (float): Average validation loss for the epoch.

  • avgValAccuracy (float): Average validation accuracy for the epoch.

Return type:

tuple

HMB.PyTorchTrainingPipeline.ImageryInferenceWithPlots(dataDir, model, modelCheckpointName=None, transform=None, useDefaultTransform=False, device=None, batchSize=1, imageSize=448, expDirs=[], overallResultsPath='Overall_Results.csv', appendResults=True, plotFontSize=16, plotFigSize=(8, 8), rocFigSize=(5, 5), dpi=720, verbose=True)[source]

Perform inference on all experiment directories and generate performance plots.

Parameters:

  • dataDir (str) – Directory containing dataset.

  • model (torch.nn.Module) – Model architecture.

  • modelCheckpointName (str, optional) – Name of the model checkpoint.

  • transform (callable, optional) – Image transform to apply.

  • useDefaultTransform (bool, optional) – Whether to use default image transform if none provided.

  • device (str or torch.device, optional) – Device to run inference on.

  • batchSize (int, optional) – Batch size for inference.

  • imageSize (int, optional) – Image size for transforms.

  • expDirs (list, optional) – List of experiment directories.

  • overallResultsPath (str, optional) – Output CSV file for overall results.

  • appendResults (bool, optional) – Whether to append to existing overall results CSV.

  • plotFontSize (int, optional) – Font size for plots.

  • plotFigSize (tuple, optional) – Figure size for confusion matrix.

  • rocFigSize (tuple, optional) – Figure size for ROC/PRC curves.

  • dpi (int, optional) – DPI for saving plots.

  • verbose (bool, optional) – Whether to print progress.

HMB.PyTorchTrainingPipeline.GenericImageryEvaluatePredictPlotSubset(datasetDir, model, subset='test', prefix='', storageDir=None, heavy=True, computeECE=True, exportFailureCases=True, eps=1e-10, saveArtifacts=True, maxSamples=None, preprocessFn=None, dpi=720)[source]

Evaluate a trained classification model on a specified dataset subset (train/val/test/all), collect predictions, compute confusion matrix and performance metrics, and optionally save predictions to a CSV file. It also generates and saves confusion matrix, ROC AUC, and PRC plots.

Parameters:

  • datasetDir (str) – Path to the dataset directory containing train/val/test splits.

  • model (callable) – A callable that takes a NumPy array (HWC, uint8 or float32) and returns a 1D array of class probabilities.

  • subset (str | None) – Dataset subset to evaluate (“train”, “val”, “test”, “all”, or None). Defaults to “test”.

  • prefix (str) – Prefix for saved figure filenames. Defaults to “”.

  • storageDir (str | None) – Directory to save predictions CSV and figures. If None, uses current directory. Defaults to None.

  • heavy (bool) – Whether to compute heavy metrics and plot ROC/PRC curves. Defaults to True.

  • computeECE (bool) – Whether to compute Expected Calibration Error (ECE). Defaults to True.

  • exportFailureCases (bool) – Whether to export misclassified samples to CSV. Defaults to True.

  • eps (float) – Small epsilon value for numerical stability in metric calculations. Defaults to 1e-10.

  • saveArtifacts (bool) – Whether to save figures and artifacts. Defaults to True.

  • maxSamples (int | None) – Maximum number of samples to evaluate. If None, evaluates all samples. Defaults to None.

  • preprocessFn (callable | None) – Optional preprocessing function to apply to each PIL image before prediction. Defaults to None.

  • dpi (int) – DPI for saved figures. Defaults to 720.

Returns:

A tuple containing:

  • str|None: Path to the saved predictions CSV file (or None when not saved).

  • dict: Computed weighted performance metrics.

  • List[int]: List of predicted class indices.

  • List[int]: List of ground truth class indices.

  • List[List[float]]: List of predicted class probabilities for each sample.

  • List[Optional[float]]: List of predicted confidences for each sample.

  • List[Dict[str, Any]]: List of prediction records for each sample.

  • List[str]: List of class names.

  • numpy.ndarray|None: Confusion matrix as a 2D numpy array, or None if not computable.

Return type:

tuple

HMB.PyTorchTrainingPipeline.GenericTabularEvaluatePredictPlotSubset(dataPath, model, targetColumn='Label', featureColumns=None, subset='test', maxRowsToRead=None, dropFirstColumn=False, prefix='', ignoreCategorical=True, tabularProcessorDir=None, numericScaler='Standard', batchSize=32, storageDir=None, heavy=True, computeECE=True, exportFailureCases=True, eps=1e-10, saveArtifacts=True, maxSamplesToEval=None, dpi=720, fontSize=13, device=None, figSize=(10, 10))[source]

Evaluate a trained classification model on a tabular dataset subset, collect predictions, compute confusion matrix and performance metrics, and optionally save predictions to a CSV file. It also generates and saves confusion matrix, ROC AUC, and PRC plots.

Parameters:

  • dataPath (str) – Path to the CSV file containing the tabular dataset.

  • model (callable) – A callable that takes a NumPy array (N, D) of features and returns a 1D array of class probabilities of shape (N, numClasses).

  • targetColumn (str) – Name of the column containing the target labels. Defaults to “Label”.

  • featureColumns (List[str] | None) – List of column names to use as features. If None, uses all columns except targetColumn. Defaults to None.

  • subset (str | None) – Dataset subset to evaluate (“train”, “val”, “test”, “all”, or None). If the CSV has a “split” column, filters by that value. Defaults to “test”.

  • maxRowsToRead (int | None) – Maximum number of rows to read from the CSV. If None, reads all rows. Defaults to None.

  • dropFirstColumn (bool) – Whether to drop the first column of the CSV (e.g., an index column). Defaults to False.

  • prefix (str) – Prefix for saved figure filenames. Defaults to “”.

  • ignoreCategorical (bool) – Whether to ignore non-numeric columns in features. Defaults to True.

  • tabularProcessorDir (str | None) – Optional directory containing tabular data processing artifacts (e.g., encoders, scalers). If provided, these will be loaded and applied to the features before prediction. Defaults to None.

  • numericScaler (str | object) – Type of numeric scaler to apply to features (“Standard”, “MinMax”, or a custom scaler object). Defaults to “Standard”.

  • batchSize (int) – Batch size for processing the dataset in batches. Defaults to 32

  • storageDir (str | None) – Directory to save predictions CSV and figures. If None, uses current directory. Defaults to None.

  • heavy (bool) – Whether to compute heavy metrics and plot ROC/PRC curves. Defaults to True.

  • computeECE (bool) – Whether to compute Expected Calibration Error (ECE). Defaults to True.

  • exportFailureCases (bool) – Whether to export misclassified samples to CSV. Defaults to True.

  • eps (float) – Small epsilon value for numerical stability in metric calculations. Defaults to 1e-10.

  • saveArtifacts (bool) – Whether to save figures and artifacts. Defaults to True.

  • maxSamplesToEval (int | None) – Maximum number of samples to evaluate. If None, evaluates all samples. Defaults to None.

  • dpi (int) – DPI for saved figures. Defaults to 720.

  • fontSize (int) – Font size for plot labels and annotations. Defaults to 13.

  • device – Optional device to run model on (e.g., “cpu” or “cuda”). If None, uses CPU. Defaults to None.

  • figSize – Tuple[int, int]: Figure size for plots. Defaults to (8, 8).

Returns:

A tuple containing:

  • str|None: Path to the saved predictions CSV file (or None when not saved).

  • dict: Computed weighted performance metrics.

  • List[int]: List of predicted class indices.

  • List[int]: List of ground truth class indices.

  • List[List[float]]: List of predicted class probabilities for each sample.

  • List[Optional[float]]: List of predicted confidences for each sample.

  • List[Dict[str, Any]]: List of prediction records for each sample.

  • List[str]: List of class names.

  • numpy.ndarray|None: Confusion matrix as a 2D numpy array, or None if not computable.

Return type:

tuple

class HMB.PyTorchTrainingPipeline.PyTorchClassificationTrainingPipeline(model, trainLoader, valLoader, allLoader, optimizer, scheduler, lossFn, noOfClasses, learningRate=0.0001, device='cuda', outputDir='Output', dpi=720, logEveryNSteps=100, useAmp=True, earlyStoppingPatience=None, judgeBy='val_loss', checkpointSaver=None, verbose=True, saveDataFrames=True, configs=None)[source]

Bases: object

Trainer helper for training, validating and evaluating a classification model using PyTorch.

This class wraps a PyTorch classification model and provides convenience wiring for the training loop, validation, checkpointing, metric computation and saving prediction artifacts. The interface mirrors other helpers in this module and centralizes common I/O paths (logs, checkpoints, predictions, and metrics) so experiments are reproducible and easy to inspect.

Parameters:

  • model (torch.nn.Module) – PyTorch classification model instance used for training and inference.

  • trainLoader (DataLoader) – Training DataLoader yielding (input, label) pairs.

  • valLoader (DataLoader) – Validation DataLoader used for periodic evaluation.

  • allLoader (DataLoader) – DataLoader covering the whole dataset (used for full predictions/exports).

  • optimizer (torch.optim.Optimizer) – Optimizer instance used to update model weights.

  • scheduler (object | None) – Learning-rate scheduler (optional) applied after each epoch/step.

  • lossFn (callable) – Loss function used for training (e.g., CrossEntropyLoss).

  • learningRate (float) – Base learning rate for bookkeeping and checkpointing (default: 1e-4).

  • device (str) – Device to run computations on (“cuda” or “cpu”).

  • outputDir (str) – Directory where logs, checkpoints and outputs will be saved.

  • dpi (int) – DPI used when saving figures (default: 720).

  • noOfClasses (int)

  • logEveryNSteps (int)

  • useAmp (bool)

  • earlyStoppingPatience (int | None)

  • judgeBy (str)

  • checkpointSaver (CheckpointSaver | None)

  • verbose (bool)

  • saveDataFrames (bool)

  • configs (Dict[str, Any] | None)

Attributes (selected):

model, device, trainLoader, valLoader, allLoader, optimizer, scheduler, lossFn, writer (SummaryWriter), checkpointDir, predsDir, metricsDir, bestMetric

Notes

  • TensorBoard summaries are written to outputDir/Logs.

  • Checkpoints are written to outputDir/Checkpoints via SaveCheckpoint wrapper.

  • Metric functions used for evaluation are mapped from HMB.PerformanceMetrics.

Initialize the PyTorch Classification Trainer.

Parameters:

  • model (torch.nn.Module) – The classification model to train and evaluate.

  • trainLoader (DataLoader) – DataLoader for training data.

  • valLoader (DataLoader) – DataLoader for validation data.

  • allLoader (DataLoader) – DataLoader covering the entire dataset for inference/export.

  • optimizer (torch.optim.Optimizer) – Optimizer for training the model.

  • scheduler – Learning rate scheduler (optional) to adjust learning rate during training.

  • lossFn – Loss function used for training (e.g., CrossEntropyLoss).

  • noOfClasses (int) – Number of output classes for the classification task.

  • learningRate (float) – Base learning rate for bookkeeping and checkpointing (default: 1e-4).

  • device (str) – Device to run computations on (“cuda” or “cpu”).

  • outputDir (str) – Directory where logs, checkpoints and outputs will be saved.

  • dpi (int) – DPI used when saving figures (default: 720).

  • logEveryNSteps (int) – Frequency of logging batch-level metrics to TensorBoard (default: 100).

  • useAmp (bool) – Whether to use automatic mixed precision for training (default: True).

  • earlyStoppingPatience (int | None) – Number of epochs with no improvement after which training will be stopped (default: None, meaning no early stopping).

  • judgeBy (str) – Metric to monitor for checkpointing and early stopping (“val_loss” or “val_accuracy”, default: “val_loss”).

  • checkpointSaver (CheckpointSaver | None) – Optional callback for saving checkpoints with custom logic (default: None).

  • verbose (bool) – Whether to print verbose logs during training and validation (default: True).

  • saveDataFrames (bool) – Whether to save train/val/test/all DataFrames to outputDir for reference (default: True).

  • configs (dict | None) – Optional configuration dictionary to save to outputDir/ConfigsUsed.json (default: None).

__init__(model, trainLoader, valLoader, allLoader, optimizer, scheduler, lossFn, noOfClasses, learningRate=0.0001, device='cuda', outputDir='Output', dpi=720, logEveryNSteps=100, useAmp=True, earlyStoppingPatience=None, judgeBy='val_loss', checkpointSaver=None, verbose=True, saveDataFrames=True, configs=None)[source]

Initialize the PyTorch Classification Trainer.

Parameters:

  • model (torch.nn.Module) – The classification model to train and evaluate.

  • trainLoader (DataLoader) – DataLoader for training data.

  • valLoader (DataLoader) – DataLoader for validation data.

  • allLoader (DataLoader) – DataLoader covering the entire dataset for inference/export.

  • optimizer (torch.optim.Optimizer) – Optimizer for training the model.

  • scheduler – Learning rate scheduler (optional) to adjust learning rate during training.

  • lossFn – Loss function used for training (e.g., CrossEntropyLoss).

  • noOfClasses (int) – Number of output classes for the classification task.

  • learningRate (float) – Base learning rate for bookkeeping and checkpointing (default: 1e-4).

  • device (str) – Device to run computations on (“cuda” or “cpu”).

  • outputDir (str) – Directory where logs, checkpoints and outputs will be saved.

  • dpi (int) – DPI used when saving figures (default: 720).

  • logEveryNSteps (int) – Frequency of logging batch-level metrics to TensorBoard (default: 100).

  • useAmp (bool) – Whether to use automatic mixed precision for training (default: True).

  • earlyStoppingPatience (int | None) – Number of epochs with no improvement after which training will be stopped (default: None, meaning no early stopping).

  • judgeBy (str) – Metric to monitor for checkpointing and early stopping (“val_loss” or “val_accuracy”, default: “val_loss”).

  • checkpointSaver (CheckpointSaver | None) – Optional callback for saving checkpoints with custom logic (default: None).

  • verbose (bool) – Whether to print verbose logs during training and validation (default: True).

  • saveDataFrames (bool) – Whether to save train/val/test/all DataFrames to outputDir for reference (default: True).

  • configs (dict | None) – Optional configuration dictionary to save to outputDir/ConfigsUsed.json (default: None).

SaveDataFrames(trainDF, valDF, testDF, allDF=None)[source]

Save train, validation, test, and all DataFrames to outputDir/Data/.

Parameters:

  • trainDF (pd.DataFrame) – Training DataFrame with features and labels.

  • valDF (pd.DataFrame) – Validation DataFrame with features and labels.

  • testDF (pd.DataFrame) – Test DataFrame with features and labels.

  • allDF (pd.DataFrame | None) – Optional combined DataFrame of all splits.

SaveCheckpoint(epoch, tag='latest')[source]
Parameters:
LoadCheckpoint(filePath, strict=True)[source]
Return type:

int

Parameters:
Train(numEpochs)[source]
Parameters:

numEpochs (int)

TrainEpoch(epoch)[source]

Run a single training epoch over the training dataset and return the average loss and accuracy.

Parameters:

epoch (int) – The current epoch number (used for logging and checkpointing).

Returns:

(avgTrainLoss, avgTrainAccuracy) for the epoch.

Return type:

tuple

Validate(epoch)[source]

Run validation over the validation set and compute loss and accuracy.

Parameters:

epoch (int) – The current epoch number (used for logging and checkpointing).

Returns:

(avgValLoss, avgValAccuracy) for the epoch.

Return type:

tuple

Inference()[source]

Run inference on the entire dataset using the allLoader and save predictions and metrics to disk.

class HMB.PyTorchTrainingPipeline.PyTorchUNetSegmentationTrainingPipeline(model, trainLoader, valLoader, allLoader, optimizer, scheduler, lossFn, learningRate=0.0001, device='cuda', outputDir='Output', dpi=720, logEveryNSteps=100, useAmp=True, earlyStoppingPatience=None, judgeBy='val_loss', checkpointSaver=None, verbose=True)[source]

Bases: object

Trainer helper for training, validating and evaluating a U-Net segmentation model using PyTorch.

This class wraps a PyTorch U-Net model and provides convenience wiring for the training loop, validation, checkpointing, metric computation and saving prediction/overlay artifacts. The interface mirrors other helpers in this module and centralizes common I/O paths (logs, checkpoints, overlays, predictions, and metrics) so experiments are reproducible and easy to inspect.

Parameters:

  • model (torch.nn.Module) – PyTorch U-Net model instance used for training and inference.

  • trainLoader (DataLoader) – Training DataLoader yielding (input, target) pairs.

  • valLoader (DataLoader) – Validation DataLoader used for periodic evaluation.

  • allLoader (DataLoader) – DataLoader covering the whole dataset (used for full predictions/exports).

  • optimizer (torch.optim.Optimizer) – Optimizer instance used to update model weights.

  • scheduler (object | None) – Learning-rate scheduler (optional) applied after each epoch/step.

  • lossFn (callable) – Loss function used for training (e.g., BCEWithLogitsLoss, DiceLoss).

  • learningRate (float) – Base learning rate for bookkeeping and checkpointing (default: 1e-4).

  • device (str) – Device to run computations on (“cuda” or “cpu”).

  • outputDir (str) – Directory where logs, checkpoints and outputs will be saved.

  • dpi (int) – DPI used when saving figures (default: 720).

  • logEveryNSteps (int)

  • useAmp (bool)

  • earlyStoppingPatience (int | None)

  • judgeBy (str)

  • checkpointSaver (CheckpointSaver | None)

  • verbose (bool)

Attributes (selected):

model, device, trainLoader, valLoader, allLoader, optimizer, scheduler, lossFn, writer (SummaryWriter), checkpointDir, overlaysDir, predsDir, actualDir, metricsDir, bestMetric

Notes

  • TensorBoard summaries are written to outputDir/Logs.

  • Checkpoints are written to outputDir/Checkpoints via SaveCheckpoint wrapper.

  • Metric functions used for evaluation are mapped from HMB.ImageSegmentationMetrics.

Initialize the PyTorch UNet Segmentation Trainer.

Parameters:

  • model (torch.nn.Module) – The U-Net model to train and evaluate.

  • trainLoader (DataLoader) – DataLoader for training data.

  • valLoader (DataLoader) – DataLoader for validation data.

  • allLoader (DataLoader) – DataLoader covering the entire dataset for inference/export.

  • optimizer (torch.optim.Optimizer) – Optimizer for training the model.

  • scheduler – Learning rate scheduler (optional) to adjust learning rate during training.

  • lossFn – Loss function used for training (e.g., BCEWithLogitsLoss, DiceLoss).

  • learningRate (float) – Base learning rate for bookkeeping and checkpointing (default: 1e-4).

  • device (str) – Device to run computations on (“cuda” or “cpu”).

  • outputDir (str) – Directory where logs, checkpoints and outputs will be saved.

  • dpi (int) – DPI used when saving figures (default: 720).

  • logEveryNSteps (int) – Frequency of logging batch-level metrics to TensorBoard (default: 100).

  • useAmp (bool) – Whether to use automatic mixed precision for training (default: True).

  • earlyStoppingPatience (int | None) – Number of epochs with no improvement after which training will be stopped (default: None, meaning no early stopping).

  • judgeBy (str) – Metric to monitor for checkpointing and early stopping (“val_loss” or “val_dice”, default: “val_loss”).

  • checkpointSaver (CheckpointSaver | None) – Optional callback for saving checkpoints with custom logic (default: None).

  • verbose (bool) – Whether to print verbose logs during training and validation (default: True).

__init__(model, trainLoader, valLoader, allLoader, optimizer, scheduler, lossFn, learningRate=0.0001, device='cuda', outputDir='Output', dpi=720, logEveryNSteps=100, useAmp=True, earlyStoppingPatience=None, judgeBy='val_loss', checkpointSaver=None, verbose=True)[source]

Initialize the PyTorch UNet Segmentation Trainer.

Parameters:

  • model (torch.nn.Module) – The U-Net model to train and evaluate.

  • trainLoader (DataLoader) – DataLoader for training data.

  • valLoader (DataLoader) – DataLoader for validation data.

  • allLoader (DataLoader) – DataLoader covering the entire dataset for inference/export.

  • optimizer (torch.optim.Optimizer) – Optimizer for training the model.

  • scheduler – Learning rate scheduler (optional) to adjust learning rate during training.

  • lossFn – Loss function used for training (e.g., BCEWithLogitsLoss, DiceLoss).

  • learningRate (float) – Base learning rate for bookkeeping and checkpointing (default: 1e-4).

  • device (str) – Device to run computations on (“cuda” or “cpu”).

  • outputDir (str) – Directory where logs, checkpoints and outputs will be saved.

  • dpi (int) – DPI used when saving figures (default: 720).

  • logEveryNSteps (int) – Frequency of logging batch-level metrics to TensorBoard (default: 100).

  • useAmp (bool) – Whether to use automatic mixed precision for training (default: True).

  • earlyStoppingPatience (int | None) – Number of epochs with no improvement after which training will be stopped (default: None, meaning no early stopping).

  • judgeBy (str) – Metric to monitor for checkpointing and early stopping (“val_loss” or “val_dice”, default: “val_loss”).

  • checkpointSaver (CheckpointSaver | None) – Optional callback for saving checkpoints with custom logic (default: None).

  • verbose (bool) – Whether to print verbose logs during training and validation (default: True).

SaveCheckpoint(epoch, tag='latest')[source]
Parameters:
LoadCheckpoint(filePath, strict=True)[source]
Return type:

int

Parameters:
Train(numEpochs)[source]
Parameters:

numEpochs (int)

PredictImage(image)[source]

Run inference on a single input image tensor and return the predicted mask tensor.

Parameters:

image (torch.Tensor) – The input image tensor to run inference on. Expected shape is [C, H, W] or [H, W].

Returns:

The predicted mask tensor. Shape will be [H, W] for single-channel output or [C, H, W] for multi-channel output.

Return type:

torch.Tensor

TrainEpoch(epoch)[source]

Run a single training epoch over the training dataset and return the average loss.

Parameters:

epoch (int) – The current epoch number (used for logging and checkpointing).

Returns:

The average training loss for the epoch.

Return type:

float

Validate(epoch)[source]

Run validation over the validation set and compute metrics.

Parameters:

epoch (int) – The current epoch number (used for logging and checkpointing).

Returns:

A dictionary containing average loss and computed metrics (e.g., MeanDice, Mean IoU, PixelAccuracy) for the validation set.

Return type:

dict

Inference()[source]

Run inference on the entire dataset using the allLoader and save predicted masks, actual masks, and overlays to disk.