Losses#

Use losses with raw network outputs during training, before calling predict. The training guide shows how to prepare predictions and targets.

B is batch size, C channels, H height and W width. Keep predictions and targets on the same device. Use integer class labels and floating targets for flow or distance maps.

Choose a loss#

Loss	Prediction and target	What it computes
`CELoss`	BCHW logits; BHW class IDs	Mean multiclass cross-entropy
`DiceLoss`	BCHW logits; BHW class IDs	Dice loss, averaged over classes present in each target image
`IoULoss`	BCHW logits; BHW class IDs	Soft IoU loss, averaged over target-present image/class pairs
`TverskyLoss`	BCHW logits; BHW class IDs	Spatial overlap with separate false-positive and false-negative coefficients
`FocalLoss`	BCHW logits; BHW class IDs	Cross-entropy that gives more weight to poorly predicted pixels
`MSE`	BCHW raw maps; matching floating targets	Mean squared error; useful for predicting flow maps
`MAE`	BCHW raw maps; matching floating targets	Mean absolute error; masking changes how errors are averaged
`BCELoss`	BCHW logits; BCHW binary/soft targets, or a shared BHW binary map	Independent binary cross-entropy per output channel
`SCELoss`	BCHW logits; BHW class IDs	Weighted forward and reverse multiclass cross-entropy
`SSIM`	Comparable floating image/maps	Compares local image patterns; can also return a contrast difference
`MSSSIM`	Matching floating BCHW maps, at least 161 pixels per axis with the default window	1 minus five-scale MS-SSIM, averaged over images and channels
`InstanSegCriterion`	B4HW maps; BHW instance IDs	Seed BCE and per-nucleus membership loss, using the model’s pixel classifier

For SSIM, set data_range=1 for normalized maps or data_range=255 for floating images in [0, 255]. The default None retains prediction-only range guessing and can change with batch contents. Floating BHW targets are shared across channels; integer BHW targets are class IDs converted to one-hot maps. The loss retains its cap at 0.5, so negative similarities have zero gradient.

For MSSSIM, supply maps in [0, 1], or set data_range=255 for floating images in [0, 255]. It does not convert logits to probabilities. Masking, weighting and smoothing options are unused.

Combine losses for a model output#

import torch

from cellseg_models_pytorch.losses import (
    CELoss,
    DiceLoss,
    JointLoss,
    MSE,
    MultiTaskLoss,
)

criterion = MultiTaskLoss({
    "cellpose": JointLoss([MSE()]),
    "type_mask": JointLoss([CELoss(), DiceLoss()]),
})
flow_prediction = torch.zeros(1, 2, 8, 8, requires_grad=True)
class_logits = torch.zeros(1, 3, 8, 8, requires_grad=True)
loss = criterion(
    yhats={"cellpose": flow_prediction, "type_mask": class_logits},
    targets={
        "cellpose": torch.zeros_like(flow_prediction),
        "type_mask": torch.zeros(1, 8, 8, dtype=torch.long),
    },
)
loss.backward()

JointLoss adds losses for one output; MultiTaskLoss adds losses across named outputs. Weights multiply each loss and need not sum to one. Supply one weight per loss. With SSIM, use return_cs=False so it returns a single tensor.

Prediction and target head keys must match and include every configured output. Supply an optional shared edgeweight map in the targets dictionary; it needs no prediction entry. Existing dummy prediction entries remain accepted. Child losses must accept the forwarded target_weight and mask; each loss controls whether and how they affect its calculation. Extra matching heads remain unused. InstanSeg supports these helpers when given the network’s pixel classifier; see InstanSeg loss.

class JointLoss(losses, weights=None)#

Sum several losses for the same prediction and target.

Choose losses and their weights.

Parameters:

losses (list[Module]) – Nonempty list of configured PyTorch loss modules. They must accept the same keyword inputs used when calling this object.
weights (list[float] | None) – One number in [0, 1] per loss. None uses 1 for every loss. Supply exactly one per loss. Weights need not sum to one.

Raises:

ValueError – losses is not a nonempty list, the weight count differs, or a supplied weight is outside [0, 1].

forward(**kwargs)#

Add the weighted losses.

Parameters:: **kwargs (object) – Passed to every loss, typically yhat, target and optional target_weight/mask tensors. Each loss must accept these inputs.
Returns:: Sum of the weighted losses. For ordinary scalar losses, the result stays on the prediction device and keeps gradients for training.
Raises:: ValueError – Loss and weight counts differ after construction.
Return type:: Tensor

class MultiTaskLoss(head_losses, loss_weights=None, **kwargs)#

Sum losses for several named model outputs.

Assign a loss to each model output.

Parameters:

head_losses (dict[str, Module]) – Dictionary mapping output names to callable loss modules, usually JointLoss objects. Names are the training dictionaries’ keys, for example cellpose and type_mask; extract tensors from network output containers before calling this loss. edgeweight is reserved for target metadata rather than a prediction head.
loss_weights (dict[str, float] | None) – Optional weights with exactly the same keys as head_losses. None uses 1 for each output. Weights need not sum to one.
**kwargs (object) – Additional constructor arguments are accepted but unused.

Raises:

ValueError – Weight keys or counts differ from head_losses, or a head uses the reserved edgeweight name.

forward(yhats, targets, mask=None, **kwargs)#

Compute and add the weighted losses for the configured outputs.

Parameters:

yhats (dict[str, Tensor]) – Dictionary of prediction tensors. Head keys must match targets and include every configured output; extra matching heads are not evaluated. A legacy dummy edgeweight entry is accepted when the real map is supplied in targets.
targets (dict[str, Tensor]) – Dictionary of matching class-label or regression tensors. An edgeweight entry supplies spatial weights to the child losses; no dummy prediction is needed. Each child must accept target_weight.
mask (Tensor | None) – Optional BHW inclusion mask passed to each output’s loss.
**kwargs (object) – Additional call arguments are accepted but unused.

Returns:

Weighted sum, normally a scalar tensor on the prediction device with gradients. With no configured outputs, returns the number 0.0.

Raises:

ValueError – Prediction/target head keys or configured weight keys disagree. A missing configured output raises KeyError before any child loss runs.

Return type:

Tensor | float

class Loss(name, apply_sd=False, apply_ls=False, apply_svls=False, edge_weight=None, class_weights=None, **kwargs)#

Select a loss by its registered name.

Use a loss class directly when you need its specific options or target format.

Choose a loss and its settings.

Parameters:

name (str) – ce, bce, dice, iou, tversky, focal, mse, mae, ssim, msssim, sce, or instanseg. InstanSeg requires the network’s probability_net as pixel_classifier and does not support the generic options below.
apply_sd (bool) – Add a squared-prediction penalty for losses supporting it.
apply_ls (bool) – Uniformly smooth targets for losses supporting it.
apply_svls (bool) – Smooth targets using neighboring pixels where supported.
edge_weight (float | None) – Use edge_weight ** target_weight as per-pixel weights where supported. None disables this weighting.
class_weights (Tensor | None) – Per-class weights for losses that support class weighting.
**kwargs (object) – Other constructor settings for the selected loss. Supported settings vary by loss; see its own entry.

Raises:

ValueError – The name is unknown. Errors from the selected loss are passed through.

forward(yhat, target, target_weight=None, **kwargs)#

Compute the selected loss.

Parameters:

yhat (Tensor) – Prediction tensor with the shape and values required by the selected loss.
target (Tensor) – Matching class labels or regression targets.
target_weight (Tensor) – Optional per-pixel exponents for edge weighting.
**kwargs – Additional forward arguments accepted by the selected loss.

Returns:

The loss on the prediction device; normally a single tensor value. SSIM can return a pair when return_cs is enabled.

Return type:

Tensor

Weighting, masking and smoothing#

Options vary by loss; check its parameters below.

apply_sd penalizes large prediction values.
apply_ls softens class labels; apply_svls also uses neighboring pixels.
class_weights weights classes when targets contain class IDs; for BCE it weights independent output channels.
edge_weight sets per-pixel weights as edge_weight ** target_weight; how these enter the calculation depends on the loss.
apply_mask=True enables masking. Masked averages differ by loss; see each loss’s description before comparing their values.

Multiclass losses#

class CELoss(apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Compute multiclass cross-entropy from class logits.

Inputs must share a device. B is batch size, C channels, H height and W width.

Compute multiclass cross-entropy from class logits.

Parameters:

apply_sd (bool) – Add 0.005 times the mean squared prediction across channels to the per-pixel loss.
apply_ls (bool) – Mix the target with a uniform class distribution.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Multiply pixel losses by the supplied mask. The average still includes all pixels, including those with zero mask weight.
edge_weight (float | None) – Multiply pixel losses by edge_weight ** target_weight. Supply that map when enabled; None disables this weighting.
class_weights (Tensor | None) – Per-class multiplier tensor of shape (C,), on the prediction device. Use with integer class labels.
**kwargs (Any) – Additional constructor arguments are accepted but unused.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute the loss for one batch.

Parameters:

yhat (Tensor) – Floating BCHW class logits.
target (Tensor) – Integer BHW class IDs in [0, C).
target_weight (Tensor | None) – Optional BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – Optional BHW inclusion weights, used only with apply_mask=True.
**kwargs (Any) – Forward options for enabled smoothing methods, such as label_smoothing, kernel_size and sigma; unused otherwise.

Returns:

A single loss value on the prediction device, with gradients for training. Combine channel errors first, then average across images and pixels. Float16 and bfloat16 logits use float32 loss and regularization arithmetic.

Return type:

Tensor

class DiceLoss(apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Compare predicted and target class regions using spatial Dice overlap.

The forward method documents which image/class pairs contribute to the mean. Inputs are raw logits, not class IDs or reconstructed instance labels.

Configure Dice loss for class masks.

Parameters:

apply_sd (bool) – Add a penalty on large prediction scores (spectral decoupling).
apply_ls (bool) – Soften target labels by sharing a small weight across all classes.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Exclude pixels with zero mask weight.
edge_weight (float | None) – Base raised to each target_weight exponent; weights spatial overlap and volume at object borders.
class_weights (Tensor | None) – Nonnegative C-element weights scaling class losses.
**kwargs (object) – Reserved for compatibility with the loss factory.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute spatial soft Dice loss from multiclass logits.

Calculate overlap separately for each image and class, then average the losses for classes present in each target image (including background). Absent target classes are excluded; their probabilities still compete with present classes through softmax. Class weights scale losses without renormalizing their mean. Mask and edge weights weight overlap and volume equally. A fully masked batch contributes zero, including regularization.

Parameters:

yhat (Tensor) – Floating BCHW logits on the same device as the target.
target (Tensor) – BHW integer class labels in [0, C).
target_weight (Tensor | None) – BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – BHW nonnegative inclusion weights when masking is enabled.
**kwargs (float | int) – Arguments for the selected target smoothing method.

Returns:

Scalar differentiable loss on the logits’ device. Half-precision logits use float32 overlap accumulation.

Raises:

ValueError – Shapes disagree or enabled edge weighting has no map.

Return type:

Tensor

class IoULoss(apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Compare predicted and target class regions using spatial soft IoU.

Configure IoU loss for class masks.

Parameters:

apply_sd (bool) – Add 0.005 times the mask-weighted mean squared logit (spectral decoupling), independently of class and edge weights.
apply_ls (bool) – Soften target labels by sharing a small weight across all classes.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Exclude pixels with zero mask weight.
edge_weight (float | None) – Base raised to each target_weight exponent; weights spatial intersection and union at object borders.
class_weights (Tensor | None) – Nonnegative C-element weights scaling class losses.
**kwargs (Any) – Reserved for compatibility with the loss factory.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute spatial soft IoU loss from multiclass logits.

Calculate soft intersection over union separately for each image and class, then average target-present image/class pairs, including background. Absent target classes are excluded; their probabilities still compete with present classes through softmax. Class weights scale losses without renormalizing their mean. Mask and edge weights apply once to both intersection and union. A fully masked batch contributes zero, including regularization.

Parameters:

yhat (Tensor) – Floating BCHW logits on the same device as the target.
target (Tensor) – BHW integer class labels in [0, C).
target_weight (Tensor | None) – BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – BHW nonnegative inclusion weights when masking is enabled.
**kwargs (Any) – Arguments for the selected target smoothing method.

Returns:

Scalar differentiable loss on the logits’ device. Half-precision logits use float32 overlap accumulation.

Raises:

ValueError – Shapes disagree or enabled edge weighting has no map.

Return type:

Tensor

Examples

>>> logits = torch.zeros(1, 2, 1, 2)
>>> target = torch.tensor([[[0, 1]]])
>>> round(IoULoss()(logits, target).item(), 6)
0.666667

class TverskyLoss(alpha=0.7, beta=0.3, apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Compare class regions with separate false-positive and false-negative weights.

Configure spatial Tversky loss for class masks.

https://arxiv.org/abs/1706.05721

Parameters:

alpha (float) – False-positive coefficient in the overlap denominator.
beta (float) – False-negative coefficient in the overlap denominator.
apply_sd (bool) – Add 0.005 times the mask-weighted mean squared logit, independently of class and edge weights (spectral decoupling).
apply_ls (bool) – Soften targets by sharing weight across classes.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Exclude pixels with zero mask weight.
edge_weight (float | None) – Base raised to each target_weight exponent; weights true positives, false positives and false negatives at object borders.
class_weights (Tensor | None) – C-element weights scaling class losses without renormalizing the mean.
**kwargs (Any) – Reserved for compatibility with the loss factory.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute spatial Tversky loss from multiclass logits.

For each image and class, use 1 - TP / (TP + alpha * FP + beta * FN), summing soft true/false positives and false negatives across height and width. Average target-present image/class pairs, including background; absent classes still compete through softmax. Presence is determined before smoothing. Mask and edge weights apply once inside all three sums. A fully masked batch has zero loss and gradients, including regularization.

Parameters:

yhat (Tensor) – Floating BCHW logits on the target’s device, before softmax.
target (Tensor) – BHW integer class IDs in [0, C).
target_weight (Tensor | None) – BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – BHW nonnegative inclusion weights when masking is enabled.
**kwargs (Any) – Options for smoothing, such as label_smoothing, kernel_size and sigma.

Returns:

Scalar differentiable loss on the logits’ device. Float16 and bfloat16 logits use float32 overlap and regularization arithmetic.

Raises:

ValueError – Shapes disagree or enabled edge weighting has no map.

Return type:

Tensor

Examples

>>> logits = torch.zeros(1, 2, 1, 2)
>>> target = torch.tensor([[[0, 1]]])
>>> round(TverskyLoss()(logits, target).item(), 6)
0.5

Tversky uses TP / (TP + alpha * FP + beta * FN) for each image and class, then averages the losses for target-present classes, including background. alpha weights false positives and beta weights false negatives. Using alpha=beta=0.5 gives Dice loss; alpha=beta=1 gives IoU loss.

class FocalLoss(alpha=0.5, gamma=2.0, apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Reduce the contribution of easy pixels to multiclass cross-entropy.

Inputs must share a device. B is batch size, C channels, H height and W width.

Reduce the contribution of easy pixels to multiclass cross-entropy.

https://arxiv.org/abs/1708.02002

Parameters:

alpha (float) – Global multiplier on the focal term.
gamma (float) – Focusing exponent; larger values down-weight easy pixels more strongly.
apply_sd (bool) – Add 0.005 times the mean squared prediction across channels to the per-pixel loss.
apply_ls (bool) – Mix the target with a uniform class distribution.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Multiply pixel losses by the supplied mask. The average still includes all pixels, including those with zero mask weight.
edge_weight (float | None) – Multiply pixel losses by edge_weight ** target_weight. Supply that map when enabled; None disables this weighting.
class_weights (Tensor | None) – Per-class multiplier tensor of shape (C,), on the prediction device. Use with integer class labels.
**kwargs (Any) – Additional constructor arguments are accepted but unused.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute the loss for one batch.

Parameters:

yhat (Tensor) – Floating BCHW class logits.
target (Tensor) – Integer BHW class IDs in [0, C).
target_weight (Tensor | None) – Optional BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – Optional BHW inclusion weights, used only with apply_mask=True.
**kwargs (Any) – Forward options for enabled smoothing methods, such as label_smoothing, kernel_size and sigma; unused otherwise.

Returns:

Return type:

Tensor

class SCELoss(alpha=0.5, beta=1.0, apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Combine forward and reverse multiclass cross-entropy.

Inputs must share a device. B is batch size, C channels, H height and W width.

Combine forward and reverse multiclass cross-entropy.

https://arxiv.org/abs/1908.06112

Parameters:

alpha (float) – Scale on the forward cross-entropy term.
beta (float) – Scale on the reverse cross-entropy term.
apply_sd (bool) – Add 0.005 times the mean squared prediction across channels to the per-pixel loss.
apply_ls (bool) – Mix the target with a uniform class distribution.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Multiply pixel losses by the supplied mask. The average still includes all pixels, including those with zero mask weight.
edge_weight (float | None) – Multiply pixel losses by edge_weight ** target_weight. Supply that map when enabled; None disables this weighting.
class_weights (Tensor | None) – Per-class multiplier tensor of shape (C,), on the prediction device. Use with integer class labels.
**kwargs (Any) – Additional constructor arguments are accepted but unused.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute the loss for one batch.

Masking multiplies per-pixel losses before the full-image mean; it does not divide by the number of included pixels. Spectral regularization is added after masking; class and edge weights scale the resulting pixel losses. Reverse cross-entropy floors targets at 1e-4 inside the logarithm, following the authors’ reference implementation, so hard-label reverse terms stay finite.

Parameters:

yhat (Tensor) – Floating BCHW class logits.
target (Tensor) – Integer BHW class IDs in [0, C).
target_weight (Tensor | None) – Optional BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – Optional BHW inclusion weights, used only with apply_mask=True.
**kwargs (Any) – Forward options for enabled smoothing methods, such as label_smoothing, kernel_size and sigma; unused otherwise.

Returns:

Return type:

Tensor

Examples

>>> logits = torch.zeros(1, 2, 1, 1)
>>> target = torch.zeros(1, 1, 1, dtype=torch.long)
>>> round(SCELoss()(logits, target).item(), 6)
4.951744

SCE floors target probabilities at 1e-4 only inside the reverse logarithm, as in the authors’ implementation. This keeps hard-label reverse terms finite without smoothing the forward term.

Binary loss#

Pass raw logits. Each channel is an independent binary task, and targets must be in [0, 1]. class_weights supplies one multiplier per channel, including when there is only one channel; it does not select background/foreground weights. Uniform label smoothing moves binary targets toward 0.5.

class BCELoss(apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Compute binary cross-entropy independently for each output channel.

Inputs must share a device. B is batch size, C channels, H height and W width.

Configure binary cross-entropy from raw logits.

Parameters:

apply_sd (bool) – Add 0.005 times the mean squared prediction across channels to the per-pixel loss.
apply_ls (bool) – Smooth binary targets toward 0.5.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Multiply pixel losses by the supplied mask. The average still includes all pixels, including those with zero mask weight.
edge_weight (float | None) – Multiply pixel losses by edge_weight ** target_weight. Supply that map when enabled; None disables this weighting.
class_weights (Tensor | None) – C-element output-channel multipliers on the prediction device, applied to both classification and regularization terms.
**kwargs (Any) – Additional constructor arguments are accepted but unused.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute the loss for one batch.

Masking multiplies per-pixel losses before the full-image mean; it does not divide by the number of included pixels. Spectral regularization is added after masking; channel and edge weights scale the resulting losses.

Parameters:

yhat (Tensor) – Floating BCHW logits, without sigmoid or clipping.
target (Tensor) – BCHW binary or soft targets in [0, 1], or a BHW binary map shared across channels. Channels are independent binary tasks, not integer multiclass labels.
target_weight (Tensor | None) – Optional BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – Optional BHW inclusion weights, used only with apply_mask=True.
**kwargs (Any) – Forward options for enabled smoothing methods, such as label_smoothing, kernel_size and sigma; unused otherwise.

Returns:

A single loss value on the prediction device, with gradients for training. Average over batch, channels and pixels, without weight normalization. Float16 and bfloat16 logits use float32 loss and regularization arithmetic.

Raises:

ValueError – Tensor shapes disagree or enabled edge weighting has no map.

Return type:

Tensor

Examples

>>> logits = torch.zeros(1, 1, 1, 2)
>>> target = torch.tensor([[[0, 1]]])
>>> round(BCELoss()(logits, target).item(), 6)
0.693147

Losses for flow and distance maps#

MSE and MAE compute half-precision inputs in float32. MSE masking retains the full-image mean; MAE divides by the mean mask weight and can also penalize predictions outside the mask. Regularization is added after masking, so a fully masked input can still contribute a penalty.

class MSE(apply_sd=False, apply_ls=False, apply_svls=False, apply_mask=False, edge_weight=None, class_weights=None, **kwargs)#

Compare raw maps using mean squared error.

Inputs must share a device. B is batch size, C channels, H height and W width.

Compare raw maps using mean squared error.

Parameters:

apply_sd (bool) – Add 0.005 times the mean squared prediction across channels to the per-pixel loss.
apply_ls (bool) – Mix the target with a uniform class distribution.
apply_svls (bool) – Smooth targets using neighboring pixels.
apply_mask (bool) – Multiply pixel losses by the supplied mask. The average still includes all pixels, including those with zero mask weight.
edge_weight (float | None) – Multiply pixel losses by edge_weight ** target_weight. Supply that map when enabled; None disables this weighting.
class_weights (Tensor | None) – Per-class multiplier tensor of shape (C,), on the prediction device. Use with integer class labels.
**kwargs (object) – Additional constructor arguments are accepted but unused.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute the loss for one batch.

Masking multiplies per-pixel losses before the full-image mean; it does not divide by the number of included pixels.

Parameters:

yhat (Tensor) – Floating BCHW regression maps; no activation is applied.
target (Tensor) – Matching BCHW maps, floating BHW maps shared across channels, or integer BHW class IDs converted to one-hot targets.
target_weight (Tensor | None) – Optional BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – Optional BHW inclusion weights, used only with apply_mask=True.
**kwargs (float | int) – Forward options for enabled smoothing methods, such as label_smoothing, kernel_size and sigma; unused otherwise.

Returns:

A single loss value on the prediction device, with gradients for training. Combine channel errors first, then average across images and pixels. Float16/bfloat16 inputs compute in float32; float64 is preserved. Regularization is added after masking and scaled by class/edge weights.

Raises:

ValueError – Shapes disagree, class weights lack integer BHW labels, or enabled edge weighting has no map.

Return type:

Tensor

Examples

>>> prediction = torch.tensor([[[[1.0, 3.0]]]])
>>> MSE()(prediction, torch.zeros_like(prediction)).item()
5.0

class MAE(alpha=0.0001, apply_sd=False, apply_mask=False, edge_weight=None, **kwargs)#

Compare raw maps using mean absolute error.

Inputs must share a device. B is batch size, C channels, H height and W width.

Compare raw maps using mean absolute error.

https://arxiv.org/pdf/1806.03535.pdf

Parameters:

alpha (float) – Penalty on absolute predictions outside the mask when masking is enabled.
apply_sd (bool) – Add 0.005 times the mean squared prediction across channels to the per-pixel loss.
apply_mask (bool) – Multiply pixel losses by the supplied mask and divide by its mean weight. alpha also controls a penalty outside the mask.
edge_weight (float | None) – Multiply pixel losses by edge_weight ** target_weight. Supply that map when enabled; None disables this weighting.
**kwargs (object) – Additional constructor arguments are accepted but unused.

forward(yhat, target, target_weight=None, mask=None, **kwargs)#

Compute the loss for one batch.

The masked sum is normalized by mean mask weight; alpha can also penalize predictions outside the mask.

Parameters:

yhat (Tensor) – Floating BCHW maps; no activation is applied.
target (Tensor) – Floating targets matching BCHW, or BHW replicated across channels.
target_weight (Tensor | None) – Optional BHW edge exponents when edge weighting is enabled.
mask (Tensor | None) – Optional BHW inclusion weights in [0, 1], used with apply_mask=True.
**kwargs (object) – Accepted but unused; smoothing and class weights are unsupported.

Returns:

A single loss value on the prediction device, with gradients for training. Combine channel errors first, then average across images and pixels. Float16/bfloat16 inputs compute in float32; float64 is preserved. Divide masked errors by mean(mask) + 1e-7. Outside-mask and spectral penalties are then added, including for fully masked inputs; edge weights scale all terms.

Raises:

ValueError – Shapes disagree or enabled edge weighting has no map.

Return type:

Tensor

Examples

>>> prediction = torch.tensor([[[[1.0, 3.0]]]])
>>> MAE()(prediction, torch.zeros_like(prediction)).item()
2.0

class SSIM(window_size=11, return_cs=False, data_range=None, **kwargs)#

Compare local image structure using SSIM dissimilarity.

Supply comparable floating maps on the same device and value scale; these losses do not convert raw class scores to probabilities.

Choose the local comparison window.

Parameters:

window_size (int) – Gaussian comparison window width, normally an odd positive integer.
return_cs (bool) – Also return the single contrast-difference value.
data_range (float | None) – Positive finite input range, normally 1 for [0, 1] maps or 255 for [0, 255] images. None keeps legacy prediction-only range guessing, which can depend on batch contents.
**kwargs (object) – Additional constructor arguments are accepted but unused.

Raises:

ValueError – The window size or explicit data range is invalid.

forward(yhat, target, **kwargs)#

Compare one batch of maps.

Parameters:

yhat (Tensor) – Floating BCHW maps. B is batch size, C channels, H height and W width.
target (Tensor) – Matching floating BCHW maps, floating BHW maps shared across channels, or integer BHW class IDs in [0, C-1] converted to one-hot maps. Inputs must share a device and value scale.
**kwargs (object) – Additional call arguments are accepted but unused.

Returns:

A scalar in [0, 0.5] computed as clamp(1 - mean_similarity, 0, 1) / 2. With return_cs=True, returns (loss, contrast_dissimilarity); the latter lies in [0, 1]. Results stay on the prediction device with gradients. Half inputs compute in float32; mixed dtypes are promoted and float64 is preserved, including under autocast.

Raises:

ValueError – Shapes disagree, inputs are empty or devices differ.
TypeError – Predictions or BCHW targets are not floating maps.
RuntimeError – Integer class IDs are outside [0, C-1].

Return type:

Tensor | tuple[Tensor, Tensor]

Examples

>>> maps = torch.zeros(1, 1, 8, 8)
>>> SSIM(data_range=1)(maps, maps).item()
0.0

class MSSSIM(window_size=11, data_range=1.0, **kwargs)#

Compare floating maps using 1 minus five-scale MS-SSIM.

Contrast is compared at each scale; brightness is compared only at the coarsest scale. Image/channel scores are combined before averaging. Inputs are not normalized or converted from logits or class labels.

Choose the Gaussian window and fixed input range.

Parameters:

window_size (int) – Positive odd Gaussian window width; sigma is 1.5.
data_range (float) – Positive finite input range, normally 1 for maps in [0, 1] or 255 for floating images in [0, 255]. Defaults to 1.
**kwargs (object) – Accepted but unused, including weighting and smoothing options.

Raises:

ValueError – The window size or data range is invalid.

forward(yhat, target, **kwargs)#

Compare one batch of maps on their existing device.

Parameters:

yhat (Tensor) – Floating BCHW maps in [0, data_range]. B is batch size, C channels, H height and W width. Each spatial axis must be at least max(16, 16 * (window_size - 1) + 1): 161 for the default window.
target (Tensor) – Matching BCHW maps with the same dtype, device and value range.
**kwargs (object) – Accepted but unused, including masks and pixel weights.

Returns:

Scalar loss in [0, 1], zero for identical maps, with gradients on the input device. Float16/bfloat16 compute in float32, including under autocast; float32/float64 retain their dtype.

Raises:

ValueError – Input shapes, devices or dtypes differ, a batch/channel axis is empty, or spatial axes are too small.
TypeError – Inputs are not floating tensors.

Return type:

Tensor

Examples

>>> maps = torch.zeros(1, 1, 161, 161)
>>> MSSSIM()(maps, maps).item()
0.0

InstanSeg loss#

Pass the network’s existing probability_net so training and reconstruction use the same learned classifier. Optimize the network parameters, which include it:

from cellseg_models_pytorch.losses import Loss
from cellseg_models_pytorch.models.instanseg import instanseg_nuclei

network = instanseg_nuclei(3, enc_name="resnet18", enc_pretrain=False)
instance_criterion = Loss(
    "instanseg", pixel_classifier=network.probability_net,
)
optimizer = torch.optim.AdamW(network.parameters(), lr=1e-4)

Apply this criterion to network(images)["nuc"].aux_map and BHW instance IDs, not cell-type labels. Train the cell-type head separately with a classification loss. Direct calls accept pred or yhat; JointLoss and MultiTaskLoss use yhat automatically.

The objective follows upstream InstanSeg: seed binary cross-entropy plus per-crop Lovász hinge for nucleus membership. Match window_size and peak_distance to reconstruction. At most 50 seed crops per image are sampled by default. Images without a foreground seed above 0.5 have only seed loss. Negative labels mark unannotated pixels; an optional binary mask excludes additional pixels from both losses and seed selection. Fully excluded images contribute zero to the batch average.

Generic smoothing, spectral decoupling, class weights and edge weights are unsupported. The mask is applied when supplied, without apply_mask=True. The old coordinate MSE is replaced by membership training; w_coord is a deprecated alias for w_inst. Reconstruction thresholds seed probabilities after sigmoid, while the input maps remain raw logits. Existing checkpoint shapes and parameter names are unchanged; the corrected objective requires retraining models trained only with the old criterion.

class InstanSegCriterion(w_seed=1.0, w_coord=None, *, pixel_classifier=None, w_inst=None, window_size=128, max_seeds=50, peak_distance=4)#

Train seed scores and per-nucleus membership with the model’s classifier.

Uses upstream InstanSeg’s seed BCE and per-crop Lovász hinge objective. Coordinate offsets and the learned feature train through membership, using the same spatial grid and classifier inputs as reconstruction.

Choose the shared classifier and contributions of the two losses.

Parameters:

w_seed (float) – Multiplier on mean binary seed cross-entropy.
w_coord (float | None) – Deprecated alias for w_inst; coordinate MSE is removed.
pixel_classifier (Module | None) – The training network’s probability_net, also used for reconstruction. Include it in the model optimizer; do not construct a separate classifier for the loss.
w_inst (float | None) – Multiplier on mean per-crop Lovász hinge. Defaults to 1.5.
window_size (int) – Maximum even crop width in pixels. Match reconstruction.
max_seeds (int) – Maximum sampled training crops per image. Random sampling bounds crop memory without a full-image mask per nucleus.
peak_distance (int) – Seed-search radius in pixels. Match reconstruction.

Raises:

ValueError – The shared classifier is missing or settings are invalid.

Example

>>> from cellseg_models_pytorch.models.instanseg import instanseg_nuclei
>>> network = instanseg_nuclei(3, enc_name="resnet18", enc_pretrain=False)
>>> criterion = InstanSegCriterion(pixel_classifier=network.probability_net)

forward(pred=None, target=None, *, yhat=None, mask=None, target_weight=None)#

Compare raw InstanSeg maps with integer nucleus-instance labels.

Parameters:

pred (Tensor | None) – Floating (B, 4, H, W) maps: XY offsets, learned shape feature, seed logits. H/W >= 2. Keep the shared classifier on this device.
target (Tensor | None) – Integer (B, H, W) labels on the same device. Zero is background, positive IDs identify nuclei, and negative labels are unannotated.
yhat (Tensor | None) – Alternative name for pred, used by the composed loss helpers. Supply exactly one of pred and yhat.
mask (Tensor | None) – Optional binary BHW inclusion mask. Excluded pixels contribute to neither loss nor seed selection, independently of their label.
target_weight (Tensor | None) – Unsupported; passing a weight map raises ValueError.

Returns:

Mean over images of w_seed times seed BCE plus w_inst times mean per-seed membership loss. Empty/no-seed images have only seed loss; fully unannotated images contribute zero. Half/bfloat16 maps use float32 loss math; float64 is preserved. Classifier precision follows its parameters with autocast disabled for the small classifier.

Raises:

ValueError – Inputs violate the shape, device, label or mask contract.

Return type:

Tensor