PyTorchHelper Module¶

Provides helper functions for working with PyTorch models and tensors.

HMB.PyTorchHelper.GetParamCount(model, restrictGrad=True, formatMillions=True)[source]¶

Get the total number of parameters in a PyTorch model, optionally restricting to only those that require gradients (trainable parameters).

Parameters:

model (torch.nn.Module) – The model to count parameters for.
restrictGrad (bool) – If True, only count parameters that require gradients (trainable). If False, count all parameters.
formatMillions (bool) – If True, format the result in millions with one decimal place. If False, return the raw count of parameters.

Returns:

Total number of parameters in the model, optionally restricted to those that require gradients.

Return type:

int

HMB.PyTorchHelper.GetPyTorchDeviceName(device)[source]¶

Get a human-readable name for a PyTorch device.

Parameters:: device (torch.device) – The device to get the name for.
Returns:: A human-readable name for the device (e.g., “CPU”, “NVIDIA GeForce RTX 3080”).
Return type:: str

HMB.PyTorchHelper.SaveModel(model, filename='model.pth')[source]¶

Save the model state to a file for later use. You can load it later using LoadModel() function from this module.

Parameters:

model (torch.nn.Module) – The model to save.
filename (str) – The name of the file to save the model to.

HMB.PyTorchHelper.SavePyTorchDict(modelDict, filename='model.pth')[source]¶

Save a PyTorch state dictionary to a file for later use. You can load it later using LoadPyTorchDict function.

Parameters:

modelDict (dict) – The state dictionary to save.
filename (str) – The name of the file to save the state dictionary to.

HMB.PyTorchHelper.LoadModel(model, filename='model.pth', device='cuda', weightsOnly=False)[source]¶

Load the model state from a file and move it to the specified device.

Parameters:

model (torch.nn.Module) – The model to load the state into.
filename (str) – The name of the file to load the model from.
device (str) – The device to load the model onto (e.g., “cpu” or “cuda”).
weightsOnly (bool) – If True, only load the weights without strict key matching.

Returns:

The model with loaded state.

Return type:

torch.nn.Module

HMB.PyTorchHelper.LoadPyTorchDict(filename='model.pth', device='cuda', weightsOnly=False)[source]¶

Load a PyTorch state dictionary from a file and map it to the specified device.

Parameters:

filename (str) – The name of the file to load the state dictionary from.
device (str) – The device to map the state dictionary onto (e.g., “cpu” or “cuda”).
weightsOnly (bool) – If True, only load the weights without strict key matching.

Returns:

The loaded state dictionary.

Return type:

dict

HMB.PyTorchHelper.SaveCheckpoint(model, optimizer, filename='chk.pth.tar', epoch=None, hparams=None)[source]¶

Save model and optimizer state to a checkpoint file. Useful for resuming training or inference later. This function saves the model’s state dictionary and the optimizer’s state dictionary to a specified file. You can load it later using LoadCheckpoint function from this module.

Parameters:

model (torch.nn.Module) – The model to save.
optimizer (torch.optim.Optimizer) – The optimizer to save.
filename (str) – The name of the file to save the checkpoint to.
epoch (int, optional) – The current epoch number to save in the checkpoint.
hparams (dict, optional) – Hyperparameters to save in the checkpoint.

HMB.PyTorchHelper.LoadCheckpoint(checkpointFile, model, optimizer, lr, device, strict=True, verbose=False)[source]¶

Load model and optimizer states from a specified checkpoint file with automatic key reconciliation.

This function initializes the model and optimizer using weights and states stored in a checkpoint file. It incorporates a robust mechanism to handle naming convention mismatches (e.g., camelCase versus snake_case) between the checkpoint keys and the model definition. If key mismatches are detected, the function employs fuzzy string matching to map unexpected keys to missing keys, validated by tensor shape compatibility to ensure weight integrity. Additionally, it updates the learning rate for the optimizer if specified.

Parameters:

checkpointFile (str) – The absolute or relative path to the checkpoint file.
model (torch.nn.Module) – The neural network model into which the state dictionary will be loaded.
optimizer (torch.optim.Optimizer) – The optimizer into which the state dictionary will be loaded. If None, optimizer state loading is skipped.
lr (float) – The learning rate to enforce for all parameter groups in the optimizer after loading.
device (torch.device) – The target device for loading tensors (e.g., torch.device(“cpu”) or “cuda”).
strict (bool) – If True, raises an error if any keys remain unmatched after reconciliation. If False, allows partial loading.
verbose (bool) – If True, it will print the handling messages.

Note

Key reconciliation uses a similarity threshold of 0.85. Weights are only mapped if the source and target tensor shapes match exactly to prevent parameter corruption.

Returns:: The loaded checkpoint dictionary containing state_dict, optimizer, and other metadata. Returns None if the checkpoint file does not exist.
Return type:: dict

HMB.PyTorchHelper.GetOptimizer(model, optimizerType='adamw', learningRate=0.0001, weightDecay=0.0001)[source]¶

Create and return a PyTorch optimizer for the given model.

Parameters:

model (torch.nn.Module) – The model whose parameters to optimize.
optimizerType (str) – The type of optimizer (“adamw”, “adam”, “sgd”, “rmsprop”, “adadelta”).
learningRate (float) – Learning rate for the optimizer.
weightDecay (float) – Weight decay (L2 penalty).

Returns:

The created optimizer.

Return type:

torch.optim.Optimizer

class HMB.PyTorchHelper.PyTorchCrossAttentionHead(hiddenSize, attentionHeadSize, dropout, bias=True, metadataDim=None)[source]¶

Bases: Module

A PyTorch module implementing a cross-attention head for transformer architectures. This module computes cross-attention between input embeddings and optional metadata embeddings.

Initialize the CrossAttentionHead module.

Parameters:

hiddenSize (int) – The dimensionality of the input embeddings.
attentionHeadSize (int) – The dimensionality of the attention head output.
dropout (float) – Dropout rate to apply to attention probabilities.
bias (bool) – Whether to include bias terms in linear projections (default: True).
metadataDim (int) – Dimensionality of the metadata input. If None, metadata projections are not created.

__init__(hiddenSize, attentionHeadSize, dropout, bias=True, metadataDim=None)[source]¶

Initialize the CrossAttentionHead module.

Parameters:

hiddenSize (int) – The dimensionality of the input embeddings.
attentionHeadSize (int) – The dimensionality of the attention head output.
dropout (float) – Dropout rate to apply to attention probabilities.
bias (bool) – Whether to include bias terms in linear projections (default: True).
metadataDim (int) – Dimensionality of the metadata input. If None, metadata projections are not created.

forward(x, metadata=None)[source]¶

Compute cross-attention between input embeddings and metadata.

Parameters:

x (torch.Tensor) – Input embeddings of shape (batchSize, seqLen, hiddenSize).
metadata (torch.Tensor) – Metadata embeddings of shape (batchSize, metadataDim).

Returns:

A tuple containing:

context (torch.Tensor): The output of the cross-attention mechanism, of shape (batchSize, seqLen, attentionHeadSize).
probs (torch.Tensor): The attention probabilities, of shape (batchSize, seqLen + 1) where the last token corresponds to metadata.

Return type:

tuple

HMB.PyTorchHelper.CreateTimmModel(modelName, numClasses, pretrained=True)[source]¶

Create a classification model using the timm library.

Parameters:

modelName (str) – Name of the model architecture to create (e.g., ‘resnet18’, ‘efficientnet_b0’).
numClasses (int) – Number of output classes for the classification task.
pretrained (bool) – If True, loads pretrained weights. Defaults to True.

Returns:

The created timm model instance.

Return type:

torch.nn.Module

HMB.PyTorchHelper.MixupFn(inputs, targets, alpha=0.4, numClasses=None)[source]¶

Apply MixUp data augmentation to inputs and targets. This function mixes pairs of samples in the batch using a Beta distribution to create new training examples. Useful for improving model generalization and robustness.

Parameters:

inputs (torch.Tensor) – Input data of shape (batch_size, …).
targets (torch.Tensor) – Target labels of shape (batch_size,) or (batch_size, num_classes).
alpha (float) – MixUp alpha parameter for Beta distribution. Default is 0.4. If alpha > 0.5, stronger mixing is applied.
numClasses (int, optional) – Number of classes for one-hot encoding if targets are indices.

Returns:

Mixed inputs and mixed targets.

mixedInputs (torch.Tensor): Mixed input data.
mixedTargets (torch.Tensor): Mixed target labels.

Return type:

tuple

HMB.PyTorchHelper.MixupCriterion(logits, softTargets)[source]¶

Compute the MixUp loss given logits and soft targets.

Parameters:

logits (torch.Tensor) – Model output logits of shape (batch_size, num_classes).
softTargets (torch.Tensor) – Soft target labels of shape (batch_size, num_classes).

Returns:

Computed MixUp loss.

Return type:

torch.Tensor

class HMB.PyTorchHelper.ExponentialMovingAverage(model=None, decay=0.9999, device=None)[source]¶

Bases: object

Implements Exponential Moving Average (EMA) for model parameters and buffers.

This class maintains a shadow copy of a model, updating its parameters via an exponential moving average of the main model’s parameters. It ensures that persistent buffers (e.g., BatchNorm running statistics) are synchronized correctly to prevent performance degradation during evaluation.

Parameters:

model (torch.nn.Module, optional) – Model to initialize EMA with. If None, EMA will be initialized on the first update.
decay (float) – Decay rate for EMA. Default is 0.9999.
device (str or torch.device, optional) – Device to store EMA weights. If None, uses the same device as the model provided during update.

update(model)[source]¶

Update EMA weights and buffers using the current model parameters.

Parameters are updated using the EMA formula. Buffers are copied directly from the current model to ensure statistics (e.g., BatchNorm) remain valid.

Parameters:: model (torch.nn.Module) – Model with current parameters to update EMA from.
Return type:: None

to(device)[source]¶

Move EMA weights to the specified device.

Parameters:: device (str or torch.device) – Device to move EMA weights to.
Return type:: None

state_dict()[source]¶

Get the state dictionary for EMA.

Returns:: State dictionary containing decay, num_updates, and module_state_dict.
Return type:: dict

load_state_dict(sd)[source]¶

Load the state dictionary for EMA.

Parameters:: sd (dict) – State dictionary to load.
Return type:: None

apply_shadow(model)[source]¶

Apply EMA weights and buffers to the given model, backing up original weights.

This method allows for evaluation using EMA weights without permanently altering the training model.

Parameters:: model (torch.nn.Module) – Model to apply EMA weights to.
Return type:: None

restore(model)[source]¶

Restore original model weights and buffers from backup.

Parameters:: model (torch.nn.Module) – Model to restore original weights to.
Return type:: None

update_from(model)[source]¶

Alias for update() method.

Parameters:: model (torch.nn.Module) – Model with current parameters to update EMA from.
Return type:: None

HMB.PyTorchHelper.ApplyDynamicQuantizationTorch(modelPath, outputPath, exampleInput=None, inputShape=None)[source]¶

Apply dynamic quantization to a PyTorch model checkpoint and save a portable TorchScript file.

This helper attempts to load a checkpoint from modelPath. It supports the following common checkpoint formats:

a dict containing a key like "model" that is a torch.nn.Module or contains a state_dict
a plain torch.nn.Module object saved directly

If the checkpoint contains only a state_dict (mapping) the architecture is not known and quantization cannot be applied by this helper; in that case the function returns None.

When successful the function will:

apply torch.quantization.quantize_dynamic to quantize Linear and Embedding modules to qint8
convert the quantized model to TorchScript (prefer torch.jit.script; fallback to tracing)
save the scripted module to outputPath

Parameters:

modelPath (str) – Path to the saved checkpoint or model file.
outputPath (str) – Path where the scripted quantized model will be saved.
exampleInput (torch.Tensor, optional) – Example input tensor to use for tracing when scripting fails.
inputShape (tuple, optional) – Alternative to exampleInput; will create a random tensor with this shape for tracing.

Returns:

Returns the outputPath on success, None on failure.

Return type:

str | None

HMB.PyTorchHelper.EvaluateModelOnPerturbations(model, run, datasetDir, storeDir, perturbations, levels, maxSamples=200, preprocessFn=None, subset='test', eps=1e-10, dpi=300)[source]¶

Evaluate a classification model under a set of input perturbations and severity levels.

This routine runs the provided model over a dataset subset (train/val/test/all) while applying controlled corruptions (noise, brightness, jpeg, occlusion, etc.) at several severity levels. It collects per-sample predictions, computes accuracy, calibration (ECE), Brier score and other auxiliary metrics, aggregates results per-perturbation and per-level, produces plots and CSV summaries, and writes a human-readable interpretation report.

Parameters:

model (callable | torch.nn.Module) – A model or a prediction callable that accepts an image (PIL or NumPy HWC) and returns a 1D array of class probabilities.
run (Any) – Identifier for the current model run used in report headers (may be a Path-like or an object with a .name attribute).
datasetDir (str | Path) – Path to the dataset root containing splits (train/val/test).
storeDir (str | Path) – Directory where results, CSVs, plots and interpretation files will be written. The directory will be created if it does not exist.
perturbations (List[str]) – List of perturbation names to evaluate. If empty, a default set of available perturbations will be used.
levels (List[float]) – List of severity levels to apply for each perturbation. If empty, default robustness levels will be used.
maxSamples (int | None) – Maximum number of samples to evaluate per run/level. If None, evaluates all available samples. Default: 200.
preprocessFn (callable | None) – Optional preprocessing function applied to each image before prediction. When provided it should accept a PIL image and return a processed image.
subset (str) – Dataset subset to evaluate: one of (“train”, “val”, “test”, “all”). Defaults to “test”.
eps (float) – Small epsilon for numerical stability in metric computations. Default: 1e-10.
dpi (int) – DPI used when saving figures. Default: 300.

Returns:

A dictionary containing the robustness evaluation results. The dictionary contains keys: such as “Dataset”, “ClassNames”, “Perturbations”, “AverageAccuracy”, “WorstAccuracy”, and an inner “RobustnessMetrics” structure with aggregated measures. In addition to the return value, the function writes these artifacts to storeDir (e.g. “RobustnessReport.json”, “RobustnessResults.csv”, “RobustnessSummary.csv”, and “Interpretation.txt”).

Return type:

dict

Example

EvaluateModelOnPerturbations(
  model=myModel,
  run=myRun,
  datasetDir="./data/cifar10",
  storeDir="./out/robustness",
  perturbations=["gaussian","jpeg"],
  levels=[0.1, 0.2, 0.3],
  maxSamples=500,
  subset="test"
)

HMB.PyTorchHelper.MeasureLatency(model, device, inputsList, runs=100, warmup=10, useCudaEvents=True)[source]¶

Measures the latency of a PyTorch model’s forward pass on a specified device, with options for warmup and precise timing.

Parameters:

model (torch.nn.Module) – The PyTorch model to evaluate.
device (torch.device) – The device on which to run the model (e.g., CPU or CUDA).
inputsList (list) – A list of input argument tuples to pass to the model’s forward method. Each tuple should contain the positional arguments for one forward pass.
runs (int) – The number of timed runs to execute for latency measurement (default: 100).
warmup (int) – The number of warmup runs to perform before timing (default: 10).
useCudaEvents (bool) – Whether to use CUDA events for timing when running on a CUDA device (default: True).

Returns:

A tuple containing the mean latency (in milliseconds), standard deviation of latency, and a dictionary of additional statistics (min, max, percentiles).

Return type:

tuple

HMB.PyTorchHelper.ComputeProfileFLOPs(model, inputsList, doCPU=True, formatGigas=True)[source]¶

Attempts to profile the FLOPs of a PyTorch model using the thop library, with an option to run on CPU for compatibility.

Parameters:

model (torch.nn.Module) – The PyTorch model to profile.
inputsList (list) – A list of input argument tuples to pass to the model’s forward method. Each tuple should contain the positional arguments for one forward pass.
doCPU (bool) – Whether to move the model and inputs to CPU for profiling (default: True). This can improve compatibility with thop if GPU profiling causes issues, but may be slower for large models.
formatGigas (bool) – If True, format the result in Gigas with one decimal place. If False, return the raw count of FLOPs.

Returns:

The calculated FLOPs for the model if profiling is successful, or None if thop is not available or if profiling fails for any reason.

Return type:

float or None

HMB.PyTorchHelper.EnableMixedPrecision(model)[source]¶

Enable automatic mixed precision training for GPU acceleration. Wraps model with AMP support and returns a gradient scaler for stable training with float16/float32 mixed precision.

Parameters:: model (nn.Module) – PyTorch model to wrap with AMP support.
Returns:: Scaler for gradient scaling during backward pass.
Return type:: torch.amp.GradScaler

class HMB.PyTorchHelper.EarlyStopping(patience=10, minDelta=0.0, mode='auto', verbose=True)[source]¶

Bases: object

Early stopping callback to halt training when validation metric plateaus.

Monitors a specified metric and stops training if no improvement is observed for a defined number of epochs.

Parameters:

patience (int) – Number of epochs to wait before stopping after no improvement.
minDelta (float) – Minimum change in monitored metric to qualify as improvement.
mode (str) – Optimization direction: “min” for loss, “max” for accuracy, “auto” to infer.
verbose (bool) – Print status messages when stopping criteria are met.

class HMB.PyTorchHelper.CheckpointSaver(savePath, saveBestOnly=True, monitor='val_loss', mode='min', verbose=True)[source]¶

Bases: object

Model checkpointing utility for saving best/latest weights.

Saves model state dictionaries to disk based on monitored metric and optional saving strategy (best-only or periodic).

Parameters:

savePath (str) – Directory path for checkpoint file storage.
saveBestOnly (bool) – Save only when monitored metric improves.
monitor (str) – Metric name to monitor for saving decisions.
mode (str) – Optimization direction: “min” for loss, “max” for accuracy.
verbose (bool) – Print messages when checkpoints are saved.

HMB.PyTorchHelper.PreparePredTensorToNumpy(predTensor, doScale2Image=False)[source]¶

Utility to convert model output tensor after the sigmoid/softmax activation to a numpy array of class indices. It can be used also with the original mask tensor if it is already in the correct format, as it handles squeezing and type conversion.

Short summary:: Takes the raw output tensor from the model (after activation) and processes it to produce a 2D numpy array of class indices. This involves squeezing unnecessary dimensions, converting boolean masks to integers if needed, and ensuring the final output is in the correct format for evaluation or visualization.

Parameters:

predTensor (torch.Tensor) – The raw output tensor from the model after activation, expected to be of shape [B, C, H, W] or [B, 1, H, W].
doScale2Image (bool) – If True, applies a threshold to convert probabilities to binary mask. Default False.

Returns:

Numpy array of shape [B, H, W] containing class indices.

Return type:

numpy.ndarray

class HMB.PyTorchHelper.PyTorchVideoTransforms(isTrain=True, mean=None, std=None, targetSize=None)[source]¶

Bases: object

Composable video transformation pipeline. This class applies a series of transformations to video data, including resizing, random cropping, horizontal flipping, and normalization. The transformations are designed to be consistent across all frames in a video sequence, ensuring that the temporal coherence of the video is maintained. The class also includes a static method for collating individual samples into batches for use with PyTorch’s DataLoader. It also provides a method to visualize the transformed video frames in a grid format, which can be useful for debugging and verification of the transformations applied to the video data.

Initialize transformation pipeline.

Parameters:

isTrain (bool) – Whether to apply training augmentations. Default is True.
mean (Optional[tuple]) – Normalization mean values per channel. Default is ImageNet mean (0.485, 0.456, 0.406).
std (Optional[tuple]) – Normalization std values per channel. Default is ImageNet std (0.229, 0.224, 0.225).
targetSize (Optional[tuple]) – Target (height, width) for resizing. Default is (224, 224).

__init__(isTrain=True, mean=None, std=None, targetSize=None)[source]¶

Initialize transformation pipeline.

Parameters:

isTrain (bool) – Whether to apply training augmentations. Default is True.
mean (Optional[tuple]) – Normalization mean values per channel. Default is ImageNet mean (0.485, 0.456, 0.406).
std (Optional[tuple]) – Normalization std values per channel. Default is ImageNet std (0.229, 0.224, 0.225).
targetSize (Optional[tuple]) – Target (height, width) for resizing. Default is (224, 224).

Return type:

None

__call__(inputVideo)[source]¶

Apply transformation pipeline to video. It takes an input video as a numpy array, applies resizing, random cropping, horizontal flipping, and normalization to each frame in the video sequence. The transformations are applied consistently across all frames to maintain temporal coherence. The output is a normalized tensor ready for input into a video model.

Parameters:: inputVideo (ndarray) – Input video array of shape (T, H, W, C).
Return type:: Tensor
Returns:: Transformed tensor of shape (C, T, H, W).

InverseTransform(inputTensor)[source]¶

Convert a normalized tensor back to a standard numpy array.

Parameters:: inputTensor (Tensor) – Normalized tensor of shape (C, T, H, W) or (B, C, T, H, W).
Return type:: ndarray
Returns:: Denormalized numpy array of shape (T, H, W, C) with values in [0, 255].

static CollateFn(sampleBatch)[source]¶

Collate function for DataLoader. It stacks the pixel value tensors and label tensors from a list of sample dictionaries into a single batched dictionary.

Parameters:: sampleBatch (list) – List of sample dictionaries.
Return type:: dict
Returns:: Batched dictionary with stacked tensors.

static VisualizeVideo(inputTensor, savePath=None)[source]¶

Display or save a grid of video frames for inspection. It takes a video tensor, converts it to a numpy array if necessary, and visualizes the frames in a grid layout using matplotlib.

Parameters:

inputTensor (Tensor) – Video tensor of shape (C, T, H, W) or numpy array (T, H, W, C).
savePath (Optional[str]) – Optional file path to save the generated visualization image.

Return type:

None