Training Octopi Models¶

Single Model Training¶

For specific use cases or when you have a known good architecture, you can train a single model directly. In this case, the command only allows for training U-Net models. To play around with more unique model configurations, or to try importing new model designs refer to the API or octopi model-explore.

Training New Models¶

This command initializes a new U-Net model using the specified architecture and training parameters.

octopi train \
    --config config.json \
    --voxel-size 10 --tomo-alg wbp \
    --tomo-batch-size 50 --val-interval 10 \
    --target-info targets,octopi,1

Fine Tuning Models¶

If we have base weights that we would like to fine-tune for new datasets, we can still use the train command. Instead of specifying the model architecture, we can simply point to the configuration file and weights to load the existing model to fine tune.

octopi train \
    --config config.json \
    --voxel-size 10 --tomo-alg wbp \
    --model-config results/model_config.yaml \
    --model-weights results/best_model_weights.pth

Model Exploration¶

Why Start with Model Exploration?¶

Rather than manually guessing which learning rates, batch sizes, or architectural choices work best for your specific tomograms, model exploration systematically tests combinations and learns from each trial to make better choices. This automated approach consistently finds better models than manual tuning.

OCTOPI's automated architecture search uses Bayesian optimization to efficiently explore hyperparameters and find optimal configurations for your specific data.

Model exploration is recommended because:

• ✅ No expertise required - Automatically finds the best model for your data
• ✅ Efficient search - Optimal performance tailored to your specific dataset
• ✅ Time savings - Avoids trial-and-error experimentation

Quick Start¶

octopi model-explore \
    --config config.json \
    --target-info targets,octopi,1 \
    --voxel-size 10 --tomo-alg denoised \
    --data-split 0.7 --model-type Unet \
    --num-trials 100 --best-metric fBeta3 \
    --study-name my-explore-job

This automatically saves results to a timestamped directory and runs 100 optimization trials by default.

Monitoring Your Training¶

Track the progress of model exploration runs in real time, inspect trial performance, and understand which hyperparameters and architectural choices drive model quality.

mlflow ui

python -m mlflow ui

INFO:     Uvicorn running on http://127.0.0.1:5000

ssh -L 5000:localhost:5000 username@login-node-hostname

mlflow.db
mlruns/

python -m mlflow ui --host 0.0.0.0 --port 5000

http://127.0.0.1:5000

#SBATCH --cpus-per-task=16
#SBATCH --mem-per-cpu=12G
#SBATCH --gpus=1
#SBATCH --constraint=[h100|h200|a100]

16 CPUs saturates the DataLoader worker pool (cap is 16). 12 GB/CPU → 192 GB total leaves comfortable headroom for worker copy-on-write drift on large caches. Drop to 8G if your cluster is tight and you cache fewer than ~15 tomograms.

Local workstation with RTX 3090/4090, A5000, etc.:

• 8–12 CPUs is enough worker throughput for most datasets.
• 64–96 GB system RAM — budget ~1–2 GB per worker plus ~0.5–2 GB per cached tomogram.

Limited to 8–16 GB VRAM, 16–32 GB system RAM:

• Lower --ncache-tomos so the training cache fits in RAM.
• Lower --batch-size from 16 → 8 to fit the GPU.
• Drop --dim-in from 96 → 64 for another ~3.4× activation-memory reduction.

Symptoms: Detected N oom_kill events in StepId=... from slurmstepd, or DataLoader worker (pid N) killed by signal: Killed after a few epochs.

In order of preference:

1. Bump --mem-per-cpu — e.g., 8G → 12G on a 16-CPU allocation adds 64 GB. Cheapest fix.
2. Lower --ncache-tomos — each cached tomogram is copy-on-write-shared to every worker; cache size × worker count is the worst-case drift.
3. Lower val_batch_size — the CopickDataModule.create default is 64; drop to 32 if validation OOMs.
4. Lower the training-worker cap — in octopi/datasets/helpers.py, adjust auto_num_workers(cap=...) from 16 → 12 → 8. Costs ~15% on epoch throughput but significantly reduces peak RAM.
5. Enable glibc trim in your SLURM script:
   Bash

   export MALLOC_TRIM_THRESHOLD_=0
   export MALLOC_ARENA_MAX=2
   

   Forces the C allocator to return freed memory to the OS more aggressively.

Symptoms: CUDA out of memory. Tried to allocate ....

In order of preference:

1. Lower --batch-size (crops per tomogram) — the biggest GPU VRAM lever. Drop 32 → 16 → 8.
2. Lower --dim-in (crop size) — 96 → 64 cuts activation memory ~3.4×.
3. Lower val_batch_size — from 128 → 64 → 32.
4. Disable EMA (use_ema=False in the API) — saves one model-sized shadow copy.
5. Shrink the model — e.g., --channels 32,64,80,80 instead of larger variants.