Training

class medtokenizers.training.Trainer(model, optimizer, accelerator, loss_fn=None, gradient_accumulation_steps=1, max_grad_norm=None, callbacks=None, channels_last=False, loss_config=None, scheduler=None, warmup_scheduler=None, warmup_steps=0, log_every_steps=1, nan_threshold=0.1, enable_compile_train_step=False, enable_gradient_checkpointing=False, disc_optimizer=None, model_ema_decay=None)

Bases: object

Accelerate-first trainer for medical tokenizers.

Supports: - Flexible callback system - Mixed precision training (via Accelerate) - Gradient accumulation - Validation during training - Checkpoint saving and loading (Accelerate directory format) - Multi-loss training with stage scheduling - NaN detection and recovery - Compound losses (VQGANLoss, VAEGANLoss) with separate discriminator optimizer - KL annealing for VAE training

Parameters:

• model (BaseTokenizer)

• optimizer (Optimizer)

• loss_fn (Optional[Module], default: None)

• gradient_accumulation_steps (int, default: 1)

• max_grad_norm (Optional[float], default: None)

• callbacks (Optional[list[Callback]], default: None)

• channels_last (bool, default: False)

• loss_config (Optional[LossConfig], default: None)

• warmup_steps (int, default: 0)

• log_every_steps (int, default: 1)

• nan_threshold (float, default: 0.1)

• enable_compile_train_step (bool, default: False)

• enable_gradient_checkpointing (bool, default: False)

• disc_optimizer (Optional[Optimizer], default: None)

• model_ema_decay (Optional[float], default: None)

__init__(model, optimizer, accelerator, loss_fn=None, gradient_accumulation_steps=1, max_grad_norm=None, callbacks=None, channels_last=False, loss_config=None, scheduler=None, warmup_scheduler=None, warmup_steps=0, log_every_steps=1, nan_threshold=0.1, enable_compile_train_step=False, enable_gradient_checkpointing=False, disc_optimizer=None, model_ema_decay=None)

Initialize trainer.

Parameters:

• model (BaseTokenizer) – Tokenizer model to train

• optimizer (Optimizer) – Optimizer for training

• accelerator – Accelerate instance (required)

• loss_fn (Optional[Module], default: None) – Loss function (defaults to Combined)

• gradient_accumulation_steps (int, default: 1) – Number of steps to accumulate gradients

• max_grad_norm (Optional[float], default: None) – Maximum gradient norm for clipping

• callbacks (Optional[list[Callback]], default: None) – List of callbacks for extensibility

• channels_last (bool, default: False) – If True, convert inputs to channels_last or channels_last_3d based on model dimensionality for better convolution throughput.

• loss_config (Optional[LossConfig], default: None) – Optional LossConfig for multi-loss training

• scheduler (default: None) – Optional main learning rate scheduler

• warmup_scheduler (default: None) – Optional warmup learning rate scheduler

• warmup_steps (int, default: 0) – Number of warmup steps (must match scheduler if provided)

• log_every_steps (int, default: 1) – Log metrics every N steps

• nan_threshold (float, default: 0.1) – Maximum allowed NaN rate (0.1 = 10%)

• enable_compile_train_step (bool, default: False) – Whether to compile the model with torch.compile

• enable_gradient_checkpointing (bool, default: False) – Whether to enable gradient checkpointing on encoder/decoder blocks (trades compute for memory)

• disc_optimizer (Optional[Optimizer], default: None) – Optional separate optimizer for discriminator parameters (used with VQGANLoss/VAEGANLoss when GAN training is enabled)

• model_ema_decay (Optional[float], default: None) – Optional EMA decay rate for model weights (e.g., 0.9999). When set, maintains an exponential moving average of model parameters. EMA weights are used for validation and saved alongside checkpoints.

init_trackers(project_name, config=None, **kwargs)

Initialize Accelerate trackers (e.g., wandb).

Parameters:

• project_name (str) – Name of project

• config (dict, default: None) – Optional configuration to log

• **kwargs – Additional arguments passed to accelerator.init_trackers

end_training()

End training and wait for all processes.

train_step(batch, sync_gradients=True)

Single training step.

Parameters:

• batch – Input batch (tensor or dict with “image” key)

• sync_gradients (bool, default: True) – If False and model supports no_sync(), skip gradient all-reduce (used for gradient accumulation intermediate steps). Default True syncs every step.

Returns:

Scalar loss value metrics: Dictionary of metrics

Return type:

loss

validate(val_loader, max_batches=None)

Run validation.

Parameters:

• val_loader (DataLoader) – Validation data loader

• max_batches (Optional[int], default: None) – Optional maximum number of batches to validate on. If None, validates on the entire dataset.

Returns:

Dictionary of validation metrics

Return type:

metrics

fit(train_loader, epochs, val_loader=None, steps_per_epoch=None, val_interval=1, resume=None, max_val_batches=None)

Train the model.

Parameters:

• train_loader (DataLoader) – Training data loader

• epochs (int) – Number of epochs to train

• val_loader (Optional[DataLoader], default: None) – Optional validation data loader

• steps_per_epoch (Optional[int], default: None) – Optional number of steps per epoch (for infinite dataloaders)

• val_interval (int, default: 1) – Run validation every N epochs

• resume (Optional[str], default: None) – Optional checkpoint directory path to resume from

• max_val_batches (Optional[int], default: None) – Optional maximum batches for validation (None = full dataset)

save_checkpoint(output_dir, metadata=None)

Save training checkpoint.

Saves Accelerate state (model, optimizer, schedulers) plus metadata.json containing epoch, global_step, warmup_steps, and scheduler state info.

Parameters:

• output_dir (str) – Directory path to save checkpoint (Accelerate format)

• metadata (dict, default: None) – Optional metadata to save (hparams, wandb_id, etc.)

Raises:

ValueError – If output_dir is not a directory path

load_checkpoint(input_dir)

Load training checkpoint.

Loads Accelerate state (model, optimizer, schedulers) and metadata.json. Validates that warmup_steps matches the current config and warns on mismatch.

Parameters:

input_dir (str) – Directory path to load checkpoint from (Accelerate format)

Raises:

ValueError – If input_dir is not a directory or is a legacy .pt file

Callbacks

class medtokenizers.training.Callback

Bases: abc.ABC

Base class for training callbacks.

Callbacks allow you to inject custom behavior at different points during training without modifying the trainer code.

on_train_begin(trainer)

Called at the beginning of training.

Parameters:

trainer (Trainer)

Return type:

None

on_train_end(trainer)

Called at the end of training.

Parameters:

trainer (Trainer)

Return type:

None

on_epoch_begin(trainer, epoch)

Called at the beginning of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

Return type:

None

on_epoch_end(trainer, epoch, metrics)

Called at the end of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

• metrics (dict)

Return type:

None

on_batch_begin(trainer, batch, batch_idx)

Called at the beginning of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

Return type:

None

on_batch_end(trainer, batch, batch_idx, loss)

Called at the end of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

• loss (float)

Return type:

None

on_validation_begin(trainer)

Called at the beginning of validation.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_end(trainer, metrics, **kwargs)

Called at the end of validation.

Parameters:

• trainer (Trainer)

• metrics (dict)

Return type:

None

class medtokenizers.training.EarlyStopping(patience=10, min_delta=0.0, monitor='val_loss')

Bases: medtokenizers.training.callbacks.Callback

Early stopping callback to stop training when validation loss stops improving.

Parameters:

• patience (int, default: 10)

• min_delta (float, default: 0.0)

• monitor (str, default: 'val_loss')

__init__(patience=10, min_delta=0.0, monitor='val_loss')

Parameters:

• patience (int, default: 10)

• min_delta (float, default: 0.0)

• monitor (str, default: 'val_loss')

on_epoch_end(trainer, epoch, metrics)

Called at the end of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

• metrics (dict)

Return type:

None

on_batch_begin(trainer, batch, batch_idx)

Called at the beginning of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

Return type:

None

on_batch_end(trainer, batch, batch_idx, loss)

Called at the end of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

• loss (float)

Return type:

None

on_epoch_begin(trainer, epoch)

Called at the beginning of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

Return type:

None

on_train_begin(trainer)

Called at the beginning of training.

Parameters:

trainer (Trainer)

Return type:

None

on_train_end(trainer)

Called at the end of training.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_begin(trainer)

Called at the beginning of validation.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_end(trainer, metrics, **kwargs)

Called at the end of validation.

Parameters:

• trainer (Trainer)

• metrics (dict)

Return type:

None

class medtokenizers.training.Checkpoint(filepath, monitor='val_loss', save_best_only=True, mode='min', verbose=True, save_every_n_epochs=None)

Bases: medtokenizers.training.callbacks.Callback

Save model checkpoints during training.

Parameters:

• filepath (str)

• monitor (str, default: 'val_loss')

• save_best_only (bool, default: True)

• mode (str, default: 'min')

• verbose (bool, default: True)

• save_every_n_epochs (Optional[int], default: None)

__init__(filepath, monitor='val_loss', save_best_only=True, mode='min', verbose=True, save_every_n_epochs=None)

Parameters:

• filepath (str)

• monitor (str, default: 'val_loss')

• save_best_only (bool, default: True)

• mode (str, default: 'min')

• verbose (bool, default: True)

• save_every_n_epochs (Optional[int], default: None)

on_epoch_end(trainer, epoch, metrics)

Called at the end of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

• metrics (dict)

Return type:

None

on_batch_begin(trainer, batch, batch_idx)

Called at the beginning of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

Return type:

None

on_batch_end(trainer, batch, batch_idx, loss)

Called at the end of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

• loss (float)

Return type:

None

on_epoch_begin(trainer, epoch)

Called at the beginning of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

Return type:

None

on_train_begin(trainer)

Called at the beginning of training.

Parameters:

trainer (Trainer)

Return type:

None

on_train_end(trainer)

Called at the end of training.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_begin(trainer)

Called at the beginning of validation.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_end(trainer, metrics, **kwargs)

Called at the end of validation.

Parameters:

• trainer (Trainer)

• metrics (dict)

Return type:

None

class medtokenizers.training.LRScheduler(scheduler)

Bases: medtokenizers.training.callbacks.Callback

Learning rate scheduling.

Parameters:

scheduler (_LRScheduler)

__init__(scheduler)

Parameters:

scheduler (_LRScheduler)

on_epoch_end(trainer, epoch, metrics)

Called at the end of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

• metrics (dict)

Return type:

None

on_batch_begin(trainer, batch, batch_idx)

Called at the beginning of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

Return type:

None

on_batch_end(trainer, batch, batch_idx, loss)

Called at the end of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

• loss (float)

Return type:

None

on_epoch_begin(trainer, epoch)

Called at the beginning of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

Return type:

None

on_train_begin(trainer)

Called at the beginning of training.

Parameters:

trainer (Trainer)

Return type:

None

on_train_end(trainer)

Called at the end of training.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_begin(trainer)

Called at the beginning of validation.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_end(trainer, metrics, **kwargs)

Called at the end of validation.

Parameters:

• trainer (Trainer)

• metrics (dict)

Return type:

None

class medtokenizers.training.Logger(print_every=10)

Bases: medtokenizers.training.callbacks.Callback

Log training progress.

Parameters:

print_every (int, default: 10)

__init__(print_every=10)

Parameters:

print_every (int, default: 10)

on_batch_end(trainer, batch, batch_idx, loss)

Called at the end of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

• loss (float)

Return type:

None

on_epoch_end(trainer, epoch, metrics)

Called at the end of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

• metrics (dict)

Return type:

None

on_batch_begin(trainer, batch, batch_idx)

Called at the beginning of each batch.

Parameters:

• trainer (Trainer)

• batch (Any)

• batch_idx (int)

Return type:

None

on_epoch_begin(trainer, epoch)

Called at the beginning of each epoch.

Parameters:

• trainer (Trainer)

• epoch (int)

Return type:

None

on_train_begin(trainer)

Called at the beginning of training.

Parameters:

trainer (Trainer)

Return type:

None

on_train_end(trainer)

Called at the end of training.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_begin(trainer)

Called at the beginning of validation.

Parameters:

trainer (Trainer)

Return type:

None

on_validation_end(trainer, metrics, **kwargs)

Called at the end of validation.

Parameters:

• trainer (Trainer)

• metrics (dict)

Return type:

None

Configuration

class medtokenizers.training.LossConfig(l1_weight=1.0, quant_weight=0.0, vgg_weight=0.0, gram_weight=0.0, laplacian_weight=0.0, stage1_epochs=0, stage1_vgg_weight=0.0, stage1_gram_weight=0.0, stage2_vgg_weight=0.0, stage2_gram_weight=0.0, stage2_warmup_epochs=0, laplacian_start_epoch=0)

Bases: object

Configuration for multi-loss training.

Supports L1 reconstruction, quantization, VGG perceptual, Gram style, and Laplacian smoothness losses with stage-based scheduling.

Parameters:

• l1_weight (float, default: 1.0)

• quant_weight (float, default: 0.0)

• vgg_weight (float, default: 0.0)

• gram_weight (float, default: 0.0)

• laplacian_weight (float, default: 0.0)

• stage1_epochs (int, default: 0)

• stage1_vgg_weight (float, default: 0.0)

• stage1_gram_weight (float, default: 0.0)

• stage2_vgg_weight (float, default: 0.0)

• stage2_gram_weight (float, default: 0.0)

• stage2_warmup_epochs (int, default: 0)

• laplacian_start_epoch (int, default: 0)

l1_weight: float = 1.0

quant_weight: float = 0.0

vgg_weight: float = 0.0

gram_weight: float = 0.0

laplacian_weight: float = 0.0

stage1_epochs: int = 0

stage1_vgg_weight: float = 0.0

stage1_gram_weight: float = 0.0

stage2_vgg_weight: float = 0.0

stage2_gram_weight: float = 0.0

stage2_warmup_epochs: int = 0

laplacian_start_epoch: int = 0

compute_stage_weights(epoch)

Compute VGG/Gram weights and current stage for given epoch.

Parameters:

epoch (int) – Current epoch number

Return type:

tuple[float, float, int]

Returns:

(vgg_weight, gram_weight, stage)

property uses_vgg: bool

Whether VGG loss is enabled in any stage.

property uses_gram: bool

Whether Gram loss is enabled in any stage.

property uses_laplacian: bool

Whether Laplacian loss is enabled.

__init__(l1_weight=1.0, quant_weight=0.0, vgg_weight=0.0, gram_weight=0.0, laplacian_weight=0.0, stage1_epochs=0, stage1_vgg_weight=0.0, stage1_gram_weight=0.0, stage2_vgg_weight=0.0, stage2_gram_weight=0.0, stage2_warmup_epochs=0, laplacian_start_epoch=0)

Parameters:

• l1_weight (float, default: 1.0)

• quant_weight (float, default: 0.0)

• vgg_weight (float, default: 0.0)

• gram_weight (float, default: 0.0)

• laplacian_weight (float, default: 0.0)

• stage1_epochs (int, default: 0)

• stage1_vgg_weight (float, default: 0.0)

• stage1_gram_weight (float, default: 0.0)

• stage2_vgg_weight (float, default: 0.0)

• stage2_gram_weight (float, default: 0.0)

• stage2_warmup_epochs (int, default: 0)

• laplacian_start_epoch (int, default: 0)