Metrics and Losses#

This module provides two categories of tools:

  • training losses for learning diffusion models

  • evaluation metrics for measuring the quality of generated samples

Training Losses#

The standard training objective for diffusion models is denoising score matching (DSM). The model is trained to recover clean data from a noisy version, with the noise scheduler handling time sampling, noise injection, and loss weighting.

MSEDSMLoss implements the MSE-based DSM loss and supports both x0-predictor and score-predictor training. WeightedMSEDSMLoss extends it with an element-wise weight tensor for masking specific spatial regions or channels (for example, land versus ocean in weather applications).

from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler
from physicsnemo.diffusion.metrics.losses import MSEDSMLoss

scheduler = EDMNoiseScheduler()

# x0-predictor training (default)
loss_fn = MSEDSMLoss(model, scheduler)

# Score-predictor training
loss_fn_score = MSEDSMLoss(
    model, scheduler,
    prediction_type="score",
    score_to_x0_fn=scheduler.score_to_x0,
)

Evaluation Metrics#

The framework provides evaluation metrics for assessing the quality of generated samples. The Fréchet Inception Distance (FID) is available using calculate_fid_from_inception_stats(), which computes the FID from precomputed Inception-v3 statistics.

API Reference#

MSEDSMLoss#

WeightedMSEDSMLoss#

calculate_fid_from_inception_stats#

physicsnemo.diffusion.metrics.fid.calculate_fid_from_inception_stats(
mu: Tensor,
sigma: Tensor,
mu_ref: Tensor,
sigma_ref: Tensor,
) → Tensor[source]#

Calculate the Fréchet Inception Distance (FID) between two sets of Inception statistics.

The Fréchet Inception Distance is a measure of the similarity between two datasets based on their Inception features (mu and sigma). It is commonly used to evaluate the quality of generated images in generative models.

Parameters:

  • mu (torch.Tensor:) – Mean of Inception statistics for the generated dataset.

  • sigma (torch.Tensor:) – Covariance matrix of Inception statistics for the generated dataset.

  • mu_ref (torch.Tensor) – Mean of Inception statistics for the reference dataset.

  • sigma_ref (torch.Tensor) – Covariance matrix of Inception statistics for the reference dataset.

Returns:

The Fréchet Inception Distance (FID) between the two datasets.

Return type:

float