Training

Training#

Training utilities and callbacks for spatial models.

Training Plans#

Training plans for spatial models.

class spatialvi.train._training_plans.SpatialTrainingPlan(module, lr=0.001, weight_decay=1e-06, spatial_weight=1.0, spatial_warmup_epochs=10, **kwargs)[source]#

Bases: TrainingPlan

Training plan with spatial-specific features.

This training plan extends the base TrainingPlan with: - Spatial loss weighting - Neighbor-aware loss computation - Multi-scale training

Parameters:

  • module (BaseModuleClass) – Neural network module.

  • lr (float) – Learning rate.

  • weight_decay (float) – Weight decay for optimizer.

  • spatial_weight (float) – Weight for spatial regularization term.

  • spatial_warmup_epochs (int) – Number of epochs to warmup spatial regularization.

  • **kwargs – Additional keyword arguments for TrainingPlan.

  • module (BaseModuleClass)

  • lr (float)

  • weight_decay (float)

  • spatial_weight (float)

  • spatial_warmup_epochs (int)

on_train_epoch_start()[source]#

Update spatial weight at epoch start.

Return type:

None

training_step(batch, batch_idx)[source]#

Training step with spatial weight.

prepare_data_per_node: bool#
allow_zero_length_dataloader_with_multiple_devices: bool#
training: bool#
class spatialvi.train._training_plans.NicheTrainingPlan(module, lr=0.001, classification_weight=1.0, niche_weight=1.0, **kwargs)[source]#

Bases: SpatialTrainingPlan

Training plan for niche-aware models.

Extends SpatialTrainingPlan with: - Niche composition loss - Classification loss for semi-supervised training - Balanced sampling for cell types

Parameters:

  • module (BaseModuleClass) – Neural network module.

  • lr (float) – Learning rate.

  • classification_weight (float) – Weight for classification loss.

  • niche_weight (float) – Weight for niche composition loss.

  • **kwargs – Additional keyword arguments.

  • module (BaseModuleClass)

  • lr (float)

  • classification_weight (float)

  • niche_weight (float)

training_step(batch, batch_idx)[source]#

Training step with niche losses.

prepare_data_per_node: bool#
allow_zero_length_dataloader_with_multiple_devices: bool#
training: bool#
class spatialvi.train._training_plans.DeconvolutionTrainingPlan(module, lr=0.001, sparsity_weight=0.1, reference_weight=0.0, **kwargs)[source]#

Bases: SpatialTrainingPlan

Training plan for spatial deconvolution models.

Extends SpatialTrainingPlan with: - Proportion sparsity regularization - Reference-guided constraints - Multi-stage training

Parameters:

  • module (BaseModuleClass) – Neural network module.

  • lr (float) – Learning rate.

  • sparsity_weight (float) – Weight for proportion sparsity.

  • reference_weight (float) – Weight for reference consistency.

  • **kwargs – Additional keyword arguments.

  • module (BaseModuleClass)

  • lr (float)

  • sparsity_weight (float)

  • reference_weight (float)

training_step(batch, batch_idx)[source]#

Training step with deconvolution losses.

configure_optimizers()[source]#

Configure optimizer with learning rate scheduling.

prepare_data_per_node: bool#
allow_zero_length_dataloader_with_multiple_devices: bool#
training: bool#

Callbacks#

Training callbacks for spatial models.

class spatialvi.train._callbacks.SpatialMetricsCallback(compute_every_n_epochs=10, metrics=None)[source]#

Bases: Callback

Callback for computing spatial metrics during training.

Parameters:

  • compute_every_n_epochs (int) – Frequency of metric computation.

  • metrics (list[str] | None) – List of metrics to compute.

  • compute_every_n_epochs (int)

  • metrics (list[str] | None)

on_train_epoch_end(trainer, pl_module)[source]#

Compute metrics at epoch end.

Return type:

None

Parameters:

  • trainer (Trainer)

  • pl_module (LightningModule)

class spatialvi.train._callbacks.NeighborSamplingCallback(neighbor_key='nn_index', include_neighbor_expr=True)[source]#

Bases: Callback

Callback for neighbor-aware batch sampling.

This callback modifies the data loader to include neighbor information in each batch.

Parameters:

  • neighbor_key (str) – Key in obsm for neighbor indices.

  • include_neighbor_expr (bool) – Whether to include neighbor expression.

  • neighbor_key (str)

  • include_neighbor_expr (bool)

on_train_batch_start(trainer, pl_module, batch, batch_idx)[source]#

Augment batch with neighbor information.

Return type:

None

Parameters:

  • trainer (Trainer)

  • pl_module (LightningModule)

  • batch (dict)

  • batch_idx (int)

class spatialvi.train._callbacks.EarlyStoppingOnSpatialLoss(patience=20, min_delta=0.0001, monitor='spatial_loss')[source]#

Bases: Callback

Early stopping based on spatial loss convergence.

Parameters:

  • patience (int) – Number of epochs to wait for improvement.

  • min_delta (float) – Minimum change to qualify as improvement.

  • monitor (str) – Metric to monitor.

  • patience (int)

  • min_delta (float)

  • monitor (str)

on_validation_epoch_end(trainer, pl_module)[source]#

Check for improvement.

Return type:

None

Parameters:

  • trainer (Trainer)

  • pl_module (LightningModule)

class spatialvi.train._callbacks.SpatialRegularizationScheduler(initial_weight=0.0, final_weight=1.0, warmup_epochs=50, schedule='linear')[source]#

Bases: Callback

Scheduler for spatial regularization weight.

Parameters:

  • initial_weight (float) – Starting weight.

  • final_weight (float) – Final weight.

  • warmup_epochs (int) – Number of epochs to warmup.

  • schedule (str) – Schedule type: “linear”, “cosine”, “step”.

  • initial_weight (float)

  • final_weight (float)

  • warmup_epochs (int)

  • schedule (str)

on_train_epoch_start(trainer, pl_module)[source]#

Update spatial weight.

Return type:

None

Parameters:

  • trainer (Trainer)

  • pl_module (LightningModule)

Contents