fruit_project.utils.trainer

Classes

Trainer

Module Contents

class fruit_project.utils.trainer.Trainer(model: torch.nn.Module, processor: transformers.AutoImageProcessor, device: torch.device, cfg: omegaconf.DictConfig, name: str, run: wandb.sdk.wandb_run.Run, train_dl: torch.utils.data.DataLoader, test_dl: torch.utils.data.DataLoader, val_dl: torch.utils.data.DataLoader)

model: torch.nn.Module

device: torch.device

scaler

cfg: omegaconf.DictConfig

optimizer: torch.optim.Optimizer

processor: transformers.AutoImageProcessor

name: str

early_stopping: fruit_project.utils.early_stop.EarlyStopping

run: wandb.sdk.wandb_run.Run

train_dl: torch.utils.data.DataLoader

test_dl: torch.utils.data.DataLoader

val_dl: torch.utils.data.DataLoader

start_epoch: int = 0

map_evaluator

accum_steps: int

scheduler: torch.optim.lr_scheduler.SequentialLR

ema: torch_ema.ExponentialMovingAverage = None

get_scheduler() → torch.optim.lr_scheduler.SequentialLR

Creates a learning rate scheduler with a warmup phase.

Returns:

The learning rate scheduler.

Return type:

SequentialLR

get_optimizer() → torch.optim.AdamW

Creates an AdamW optimizer with a differential learning rate for the backbone and the rest of the model (head), following standard fine-tuning practices.

Returns:

The configured optimizer.

Return type:

AdamW

move_labels_to_device(batch: transformers.BatchEncoding) → transformers.BatchEncoding

Moves label tensors within a batch to the specified device.

Parameters:

batch (BatchEncoding) – The batch containing labels.

Returns:

The batch with labels moved to the device.

Return type:

BatchEncoding

nested_to_cpu(obj: Any) → Any

Recursively moves tensors in a nested structure (dict, list, tuple) to CPU.

Parameters:

obj – The object containing tensors to move.

Returns:

The object with all tensors moved to CPU.

format_targets_for_cm(targets: List[Dict]) → List[Dict]

Formats raw targets for torchmetrics and confusion matrix. This is a helper for the confusion matrix, as MAPEvaluator handles its own formatting.

train(current_epoch: int) → Dict[str, float]

Performs one epoch of training.

Parameters:

current_epoch (int) – The current epoch number.

Returns:

The average training loss for the epoch.

Return type:

float

eval(val_dl: torch.utils.data.DataLoader, current_epoch: int, calc_cm: bool = False) → Tuple[dict[str, float], dict[str, Any], fruit_project.utils.metrics.ConfusionMatrix | None]

_run_eval(val_dl: torch.utils.data.DataLoader, current_epoch: int, calc_cm: bool = False) → Tuple[dict[str, float], dict[str, Any], fruit_project.utils.metrics.ConfusionMatrix | None]

Evaluates the model on a given dataloader.

Parameters:

• test_dl (DataLoader) – The dataloader for evaluation.

• current_epoch (int) – The current epoch number.

• calc_cm (bool, optional) – Whether to calculate and return a confusion matrix. Defaults to False.

Returns:

A tuple containing:

• loss (Dict): The evaluation loss.

• test_map (float): The mAP@.5-.95.

• test_map50 (float): The mAP@.50.

• test_map_50_per_class (torch.Tensor): The mAP@.50 for each class.

• cm (ConfusionMatrix | None): The confusion matrix if calc_cm is True, else None.

Return type:

tuple

fit() → None

Runs the main training loop for the specified number of epochs.

_build_save_dict(epoch)

_save_checkpoint(epoch: int) → str

Saves a checkpoint of the model, optimizer, scheduler, and scaler states.

Parameters:

epoch (int) – The current epoch number.

Returns:

The path to the saved checkpoint file.

Return type:

str

_load_checkpoint(path: str, model_only: bool = True) → None

Loads a checkpoint and restores the state of the model, optimizer, scheduler, and scaler.

Parameters:

path (str) – The path to the checkpoint file.