Adding New Models¶

This page covers how to integrate custom neural network architectures into octopi for both direct training and automated model exploration.

Overview¶

Octopi supports custom model architectures through a standardized template system. By following the model template structure, you can seamlessly integrate any PyTorch-based architecture into octopi's training and automated hyperparameter search workflows.

Custom models enable you to:

Experiment with novel architectures tailored to your specific particle types
Integrate state-of-the-art models from recent research papers
Leverage automated hyperparameter search for your custom architectures
Maintain compatibility with octopi's training and inference pipelines

Model Template Structure¶

All custom models must inherit from the standardized template located at octopi/models/ModelTemplate.py. The template defines three essential methods that enable full integration with octopi's workflows.

Required Methods¶

build_model(config): Constructs the model from a configuration dictionary
bayesian_search(trial, num_classes): Defines the hyperparameter search space for automated optimization
get_model_parameters(): Returns the stored model configuration

Basic Template¶

Python

import torch.nn as nn
import torch

class MyCustomModel:
    def __init__(self):
        """
        Initialize the model template.
        """
        # Placeholder for the model and config
        self.model = None
        self.config = None

    def build_model(self, config: dict):
        """
        Build the model based on provided parameters in a config dictionary.

        Args:
            config (dict): Configuration dictionary containing model parameters

        Returns:
            torch.nn.Module: The constructed model
        """
        self.config = config

        # Your model construction logic here
        # self.model = YourModelClass(**config)

        return self.model

    def bayesian_search(self, trial, num_classes: int):
        """
        Define the hyperparameter search space for Bayesian optimization.

        Args:
            trial (optuna.trial.Trial): Optuna trial object for suggesting parameters
            num_classes (int): Number of classes in the dataset

        Returns:
            torch.nn.Module: The constructed model with suggested parameters
        """
        # Define search space using trial.suggest_* methods
        # param1 = trial.suggest_int("param1", min_val, max_val)
        # param2 = trial.suggest_categorical("param2", [choice1, choice2])

        # Create config dictionary
        self.config = {
            'architecture': 'MyCustomModel',
            'num_classes': num_classes,
            # Add your suggested parameters here
        }

        return self.build_model(self.config)

    def get_model_parameters(self):
        """Retrieve stored model parameters."""
        if self.model is None:
            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")

        return self.config

Complete Example: Custom U-Net¶

Here's a complete example showing how to integrate a custom U-Net architecture:

Python

from monai.networks.nets import UNet
import torch.nn as nn
import torch

class MyUNet:
    def __init__(self):
        # Placeholder for the model and config
        self.model = None
        self.config = None

    def build_model(self, config: dict):
        """Creates the U-Net model based on provided parameters."""

        self.config = config
        self.model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=config['num_classes'],
            channels=config['channels'],
            strides=config['strides'],
            num_res_units=config['num_res_units'],
            dropout=config['dropout']
        )
        return self.model

    def bayesian_search(self, trial, num_classes: int):
        """Defines the Bayesian optimization search space and builds the model."""

        # Define the search space
        num_layers = trial.suggest_int("num_layers", 3, 5)
        hidden_layers = trial.suggest_int("hidden_layers", 1, 3)
        base_channel = trial.suggest_categorical("base_channel", [8, 16, 32])
        num_res_units = trial.suggest_int("num_res_units", 1, 3)
        dropout = trial.suggest_float("dropout", 0.0, 0.3)

        # Create channel sizes and strides
        downsampling_channels = [base_channel * (2 ** i) for i in range(num_layers)]
        hidden_channels = [downsampling_channels[-1]] * hidden_layers
        channels = downsampling_channels + hidden_channels
        strides = [2] * (num_layers - 1) + [1] * hidden_layers

        # Create config dictionary
        self.config = {
            'architecture': 'MyUNet',
            'num_classes': num_classes,
            'channels': channels,
            'strides': strides,
            'num_res_units': num_res_units,
            'dropout': dropout
        }

        return self.build_model(self.config)

    def get_model_parameters(self):
        """Retrieve stored model parameters."""
        if self.model is None:
            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")

        return self.config

Integration Steps¶

Step 1: Create Your Model Class¶

Create a new Python file in the octopi/models/ directory (e.g., MyCustomModel.py) following the template structure shown above.

Step 2: Register Your Model¶

Add your model to octopi's model registry by adding an elif branch to get_model() in octopi/models/common.py:

Python

# In octopi/models/common.py
from octopi.models import (
    Unet, SwinUNETR, DynUnet, SegResNet,
    MyCustomModel  # import your new module
)

def get_model(architecture):
    if architecture == "Unet":
        model = Unet.myUNet()
    elif architecture == "SwinUNETR":
        model = SwinUNETR.mySwinUNETR()
    elif architecture == "DynUNet":
        model = DynUnet.myDynUNet()
    elif architecture == "SegResNet":
        model = SegResNet.mySegResNet()
    elif architecture == "MyCustomModel":          # add this
        model = MyCustomModel.MyCustomModel()      # add this
    else:
        raise ValueError(f"Model type {architecture} not supported!")
    return model

Currently shipped architectures

Octopi currently ships Unet, SwinUNETR, DynUNet, and SegResNet in octopi/models/. Older revisions of this page referenced AttentionUnet and MedNeXt — those were never merged; use this template to add them yourself if needed.

Step 3: Use in Training¶

Once registered, your model can be used in training workflows:

Python

from octopi.datasets.config import DataGeneratorConfig
from octopi.workflows import train

# Build the data generator (see the Training guide for details)
cfg = DataGeneratorConfig(
    config=config,
    name='targets', user_id='octopi', session_id='1',
    tomo_uris='denoised@10.0',
)
data_generator = cfg.create_data_generator()

# Model configuration for your custom architecture
model_config = {
    'architecture': 'MyCustomModel',
    'num_classes': data_generator.Nclasses,
    'your_param1': value1,
    'your_param2': value2,
    # Add your specific parameters
}

# Train using your custom model
train(
    data_generator, loss_function,
    model_config=model_config,
    model_save_path='results',
)

Step 4: Use in Model Exploration¶

Architecture search runs through the octopi model-explore CLI — there is no separate Python entry point for it. --model-type is a fixed click.Choice, so also add your model's name to that choice list in octopi/entry_points/run_optuna.py alongside the get_model() registration from Step 2:

Bash

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