Training¶

This page covers data preparation and training deep learning models for 3D particle picking using octopi.

Target Creation¶

Before training, you need to create training targets from existing particle annotations. We can explicitly define all the objects we'd like to query with a single ordered list of tuples (object_name, user_id, session_id) defining source annotations. Each entry can either be point coordinates for protein locations, or a continuous segmentation such as for membranes or organelles generated from software such as membrain-seg or saber — the type is auto-detected per object from the CoPick config, so both kinds are listed together in one list.

The target creation process uses the following key parameters:

• targets: Ordered list of tuples (object_name, user_id, session_id) defining source annotations — particle picks and continuous segmentations can be freely mixed; labels are assigned sequentially in list order
• radius_scale: Scale factor for creating spherical targets relative to object radius defined in config
• run_ids: Optional subset of tomograms (None for all available)

copick browse -c config.json

copick browse -ds datasetID

Creating Training Targets¶

from octopi.entry_points.run_create_targets import create_sub_train_targets

# Configuration
config = '../config.json'
target_name = 'targets'
target_user_id = 'octopi'
target_session_id = '1'

# Tomogram parameters
voxel_size = 10.012
tomogram_algorithm = 'wbp-denoised-denoiset-ctfdeconv'
radius_scale = 0.7  # Fraction of particle radius for sphere targets

# Define source annotations — particle picks and continuous segmentations, in the
# order their labels should be assigned (type is auto-detected from the config)
targets = [
    ('ribosome', 'data-portal', '1'),
    ('virus-like-particle', 'data-portal', '1'),
    ('apoferritin', 'data-portal', '1'),
    ('membrane', 'membrain-seg', '2'),  # continuous segmentation target
]

# Create targets
create_sub_train_targets(
    config, targets, voxel_size, radius_scale, 
    tomogram_algorithm, target_name, target_user_id, target_session_id, None
)

Model Training¶

Once training targets are created, you can train deep learning models for particle segmentation. The training process requires defining data splits, model architecture, loss functions, and training parameters. Key training parameters include:

• target_info: Tuple specifying which training targets to use (name, user_id, session_id)
• trainRunIDs/validateRunIDs: Optional lists defining which tomograms to use for training vs validation. If None, octopi automatically uses all runs that have the desired segmentation target.
• model_config: Dictionary specifying the neural network architecture and parameters. We have a few default models to chose from, including the opportunity to import new model architectures.
• loss_function: Loss function appropriate for segmentation tasks (handles class imbalance)
• use_ema: Whether to use exponential moving average for more stable training

Basic Training Setup¶

from octopi.datasets.config import DataGeneratorConfig
from octopi.workflows import train
from monai.losses import TverskyLoss

config = 'train_config.json'
results_folder = 'model_output'

# Build the data generator
cfg = DataGeneratorConfig(
    config=config,
    name='targets', user_id='octopi', session_id='1',
    tomo_uris='denoised@10.012',
    trainRunIDs=['run_001', 'run_002', 'run_003'],
    validateRunIDs=['run_004'],
)
data_generator = cfg.create_data_generator()

# Model architecture
model_config = {
    'architecture': 'Unet',
    'num_classes': data_generator.Nclasses,  # objects + background
    'dim_in': 80,
    'strides': [2, 2, 1],
    'channels': [48, 64, 80, 80],
    'dropout': 0.0,
    'num_res_units': 1,
}

# Loss function
loss_function = TverskyLoss(
    include_background=True,
    to_onehot_y=True,
    softmax=True,
    alpha=0.3,
    beta=0.7
)

# Train the model
train(
    data_generator, loss_function,
    model_config=model_config,
    model_save_path=results_folder
)

Cross Validation¶

Cross-validation is essential for robust model evaluation, especially when working with limited tomography data. It helps assess how well your model generalizes to unseen tomograms and prevents overfitting to specific data characteristics. This approach systematically trains multiple models, each time holding out a different tomogram for validation.

# Cross-validation training
def train_cross_validation(split_id):
    runIDs = get_all_run_ids()
    valRunID = [runIDs[split_id]]
    trainRunIDs = runIDs[:split_id] + runIDs[split_id+1:]

    results_folder = f'split_{split_id}'

    train(
        ...
        trainRunIDs=trainRunIDs,
        validateRunIDs=valRunID,
        ...
    )

Loss Functions¶

Choosing the right loss function is crucial for effective particle segmentation, especially given the class imbalance typical in cryo-ET data where background voxels vastly outnumber particle voxels.

MONAI Loss Functions:

from monai.losses import FocalLoss, TverskyLoss, DiceLoss, GeneralizedDiceLoss

• TverskyLoss: Generalizes Dice loss with separate weights for false positives (alpha) and false negatives (beta). Excellent for imbalanced data - use higher beta (e.g., 0.7) to penalize missed particles more than false detections.
• DiceLoss: Standard choice for segmentation, works well when classes are relatively balanced.

Octopi Custom Loss Functions:

from octopi.utils.losses import FocalTverskyLoss, WeightedFocalTverskyLoss

• FocalTverskyLoss: Combines Tversky's class balancing with Focal's hard example mining. The gamma parameter adds focusing power to the Tversky formulation.
• WeightedFocalTverskyLoss: Weighted combination of separate Focal and Tversky losses, allowing fine-tuned control over both class imbalance and hard example emphasis through weight_tversky and weight_focal parameters.

Model Exploration¶

In cases where we'd like to automatically explore the model architecture landscape to determine which model configuration would be optimal for our given experiment, octopi uses Optuna's Bayesian optimization to explore various loss functions and their associated hyperparameters, as well as architecture parameters defined in each model class. This automated approach significantly reduces the number of input parameters you need to provide, as the system intelligently searches through the hyperparameter space.

The automated search process efficiently explores combinations of:

• Loss function types and their hyperparameters (alpha, beta, gamma values)
• Model architecture parameters (channel sizes, dropout rates, number of residual units)
• Model family — Unet, SwinUNETR, DynUNet, or SegResNet

This is particularly valuable when working with new datasets or particle types where optimal configurations are unknown.

octopi model-explore \
    --config config.json \
    --tomo-uri denoised@10.0 \
    --target-uri targets:octopi/1 \
    --model-type Unet \
    --num-trials 100

Next Steps¶

Once you've successfully trained your model, you're ready to move to the inference stage. The trained model outputs (.pth weights and .yaml configuration) from this training process will be used directly in the inference pipeline.

Ready to apply your trained model? Head to the Inference Guide to learn how to:

• Run segmentation on new tomograms using your trained weights
• Perform particle localization from segmentation masks
• Evaluate your model's performance against ground truth data

Want to experiment with custom architectures? Check out the Adding New Models template to:

• Implement custom neural network architectures
• Integrate new model types into the octopi framework
• Extend the automated model search capabilities

Pro tip: If you're unsure about optimal hyperparameters, consider running the automated model exploration first before manual training - it can save significant time and often discovers configurations you might not have considered!