{"id":51834332,"url":"https://github.com/cell-observatory/cell_observatory_platform","last_synced_at":"2026-07-22T19:03:29.210Z","repository":{"id":279858484,"uuid":"938353036","full_name":"cell-observatory/cell_observatory_platform","owner":"cell-observatory","description":"Training backend for Cell Observatory models ","archived":false,"fork":false,"pushed_at":"2026-07-21T19:19:52.000Z","size":3158,"stargazers_count":4,"open_issues_count":1,"forks_count":2,"subscribers_count":3,"default_branch":"develop","last_synced_at":"2026-07-21T21:25:15.942Z","etag":null,"topics":["3d","4d","5d","jepa","machine-learning","mae","masked-autoencoder","microscopy","pretraining","self-supervised-learning"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cell-observatory.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-02-24T20:24:41.000Z","updated_at":"2026-03-27T09:35:44.000Z","dependencies_parsed_at":"2025-02-28T05:31:18.599Z","dependency_job_id":"431092f4-ab9e-4ce5-b3af-9fd5b5ec6802","html_url":"https://github.com/cell-observatory/cell_observatory_platform","commit_stats":null,"previous_names":["cell-observatory/platform","cell-observatory/cell_observatory_platform"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/cell-observatory/cell_observatory_platform","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cell-observatory%2Fcell_observatory_platform","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cell-observatory%2Fcell_observatory_platform/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cell-observatory%2Fcell_observatory_platform/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cell-observatory%2Fcell_observatory_platform/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cell-observatory","download_url":"https://codeload.github.com/cell-observatory/cell_observatory_platform/tar.gz/refs/heads/develop","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cell-observatory%2Fcell_observatory_platform/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":35773465,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-07-20T02:08:10.276Z","status":"online","status_checked_at":"2026-07-22T02:00:06.236Z","response_time":124,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["3d","4d","5d","jepa","machine-learning","mae","masked-autoencoder","microscopy","pretraining","self-supervised-learning"],"created_at":"2026-07-22T19:03:28.480Z","updated_at":"2026-07-22T19:03:29.191Z","avatar_url":"https://github.com/cell-observatory.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\n# Cell Observatory Platform\n**The Cell Observatory Platform** is a comprehensive framework for training and evaluating machine learning models on biological image and video datasets. Built with [PyTorch](https://pytorch.org/), orchestrated and scaled with [Ray](https://www.ray.io/), model sharding via [DeepSpeed](https://www.deepspeed.ai/) or native PyTorch parallelism ([TorchTitan](https://github.com/pytorch/torchtitan)), and flexibly configured using [Hydra](https://hydra.cc/), it provides a modular architecture for easy customization and extension.\n\n- [Cell Observatory Platform](#cell-observatory-platform)\n- [Installation](#installation)\n  - [Docker images](#docker-images)\n  - [Clone the repository to your host system](#clone-the-repository-to-your-host-system)\n  - [Setup Supabase and W\\\u0026B accounts](#setup-supabase-and-wb-accounts)\n  - [Setup environment variables](#setup-environment-variables)\n- [Running docker image](#running-docker-image)\n  - [Running docker image on a cluster via apptainer](#running-docker-image-on-a-cluster-via-apptainer)\n    - [amd64/x86\\_64](#amd64x86_64)\n    - [arm64/aarch64](#arm64aarch64)\n  - [Building a new apptainer image with a different torch version](#building-a-new-apptainer-image-with-a-different-torch-version)\n- [Get started](#get-started)\n  - [Local setup](#local-setup)\n  - [Cluster setup](#cluster-setup)\n- [Architecture Overview](#architecture-overview)\n- [Data Pipeline](#data-pipeline)\n  - [Databases](#databases)\n  - [Ray Dataloader](#ray-dataloader)\n  - [Preprocessors](#preprocessors)\n  - [Transforms](#transforms)\n  - [MaskGenerator](#maskgenerator)\n- [Models](#models)\n  - [Pretraining Models](#pretraining-models)\n  - [Detection \\\u0026 Segmentation](#detection--segmentation)\n  - [Backbones](#backbones)\n  - [Layers](#layers)\n- [Training](#training)\n  - [EpochBasedTrainer (DeepSpeed)](#epochbasedtrainer-deepspeed)\n  - [ParallelEpochBasedTrainer (TorchTitan)](#parallelepochbasedtrainer-torchtitan)\n  - [Hooks](#hooks)\n- [Logging \\\u0026 Experiment Tracking](#logging--experiment-tracking)\n  - [Event Recording](#event-recording)\n  - [W\\\u0026B Integration](#wb-integration)\n  - [Metrics Processing](#metrics-processing)\n- [Inference](#inference)\n- [Evaluation](#evaluation)\n- [Profiling](#profiling)\n- [Configuration layout](#configuration-layout)\n- [License](#license)\n\n\n# Installation\n\n## Docker [images](https://github.com/cell-observatory/cell_observatory_platform/pkgs/container/cell_observatory_platform)\nOur prebuilt image with Python, Torch, and all packages installed for you.\n```shell\ndocker pull ghcr.io/cell-observatory/cell_observatory_platform:develop_torch_26_01\n```\n\n## Clone the repository to your host system\n```shell\ngit clone --recurse-submodules https://github.com/cell-observatory/cell_observatory_platform.git\n```\n\nTo later update to the latest, greatest.\n```shell\ngit pull --recurse-submodules\n```\n\n\u003e [!NOTE]\n\u003e If you want to run a local version of the image, see the [Dockerfile](https://github.com/cell-observatory/cell_observatory_platform/blob/main/Dockerfile)\n\n## Setup Supabase and W\u0026B accounts\n\nYou will need to create a Supabase and W\u0026B account to use the platform.\nSupabase can be found at [Cell Observatory Database](https://supabase.com/dashboard/org/yrgvnbckfmhfyxgzzkqb), and W\u0026B can be found at [Cell Observatory Dashboard](https://wandb.ai/cell-observatory).\n\nOnce you have created your Supabase and W\u0026B accounts, you'll need to add your API keys in the environment variables as described below.\n\n## Setup environment variables\nRename `.env.example` file to `.env` which will be automatically loaded into the container and will be gitignored. The Supabase related environment variables enable database functionality. The W\u0026B API key enables logging functionality. The `REPO_NAME`, `DATA_DIR`, and `STORAGE_SERVER_DIR` environment variables are leverged in the `configs/paths` configuration files to ensure that jobs run and save outputs as expected.\n\n\u003e [!NOTE]\n\u003e `STORAGE_SERVER_DIR` is usually set to the root directory of your files. See example below:\n\u003e ```shell\n\u003e STORAGE_SERVER_DIR=\"/clusterfs/scratch/user/\"\n\u003e REPO_DIR=\"/clusterfs/scratch/user/cell_observatory_platform\"\n\u003e DATA_DIR=\"/clusterfs/scratch/user/cell_observatory_data\"\n\u003e ```\n\n```shell\nSUPABASE_USER=REPLACE_ME_WITH_YOUR_SUPABASE_USERNAME\nSUPABASE_PASS=REPLACE_ME_WITH_YOUR_SUPABASE_PASSWORD\nTRINO_USER=REPLACE_ME_WITH_YOUR_TRINO_USERNAME\nTRINO_PASS=REPLACE_ME_WITH_YOUR_TRINO_PASSWORD\nSUPABASE_STAGING_ID=REPLACE_ME_WITH_YOUR_SUPABASE_STAGING_ID\nSUPABASE_PROD_ID=REPLACE_ME_WITH_YOUR_SUPABASE_PROD_ID\nWANDB_API_KEY=REPLACE_ME_WITH_YOUR_WANDB_API_KEY\n\nSUPABASE_STAGING_URI=\"postgresql://${SUPABASE_USER}.${SUPABASE_STAGING_ID}:${SUPABASE_PASS}@aws-0-us-east-1.pooler.supabase.com:5432/postgres\"\nSUPABASE_PROD_URI=\"postgresql://${SUPABASE_USER}.${SUPABASE_PROD_ID}:${SUPABASE_PASS}@aws-0-us-east-1.pooler.supabase.com:5432/postgres\"\n\nREPO_NAME=cell_observatory_platform # TODO: replace with your repo name if you renamed it\nREPO_DIR=REPLACE_ME_WITH_YOUR_ROOT_REPO_DIR\nDATA_DIR=REPLACE_ME_WITH_YOUR_ROOT_DATA_DIR_WHERE_DATA_WILL_BE_SAVED\nSTORAGE_SERVER_DIR=REPLACE_ME_WITH_YOUR_STORAGE_SERVER_DIR_WHERE_DATA_SERVER_IS_MOUNTED\nPYTHONPATH=REPLACE_ME_WITH_YOUR_ROOT_REPO_DIR\n````\n\n\u003e [!IMPORTANT]\n\u003e Username/password and IDs for supabase will be provided upon request. \n\n\n# Running docker image\n\nTo run docker image, cd to repo directory or replace `$(pwd)` with your local path for the repository.\n```shell\ndocker run --network host -u 1000 --privileged -v $(pwd):/workspace/cell_observatory_platform -w /workspace/cell_observatory_platform --env PYTHONUNBUFFERED=1 --pull missing -it --rm  --ipc host --gpus all ghcr.io/cell-observatory/cell_observatory_platform:develop_torch_26_01 bash\n```\n\n## Running docker image on a cluster via apptainer\nRunning an image on a cluster typically requires an apptainer version of the image, which can be generated by:\n\n### amd64/x86_64\n```shell\napptainer pull --arch amd64 --force develop_torch_26_01.sif docker://ghcr.io/cell-observatory/cell_observatory_platform:develop_torch_26_01\n```\n\n### arm64/aarch64\n```shell\napptainer pull --arch arm64 --force develop_torch_26_01_arm64.sif docker://ghcr.io/cell-observatory/cell_observatory_platform:develop_torch_26_01\n```\n\n## Building a new apptainer image with a different torch version \nFirst, you need to build an apptainer image for torch from the containers provided by Nvidia (e.g., `26.01-py3` from this [catalog](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags)):\n``` shell\napptainer pull --arch arm64 --force pytorch_26.01-py3_arm64.sif docker://nvcr.io/nvidia/pytorch:26.01-py3\n```\n\nThen you can run the following command to build a complete image:  \n``` shell\napptainer build --arch amd64 --nv --force develop_torch_26_01.sif apptainerfile.def\n```\n\n\u003e [!IMPORTANT]\n\u003e Make sure to pass in the right argument for your system  (`amd64` or `arm64`)\n\n# Get started\n\nAll jobs are orchestrated on top of a **Ray cluster** and launched through our `manager.py` script, which facilitates cluster resource allocation and Ray cluster setup. You may decide whether to run jobs locally or on a cluster by setting the `launcher_type` variable in `configs/clusters/*.yaml`. We show how to run jobs locally and on SLURM or LSF clusters below.\n\n## Local setup\n\nExample job configs are located in the `configs/experiments` folder. For local jobs, you can use our existing `configs/paths/local.yaml` and `configs/clusters/local.yaml` configurations.\n\n### 1. Update experiment name\n```yaml\nexperiment_name: test_cell_observatory_platform\nwandb_project: test_cell_observatory_platform\n```\n\n### 2. Update your paths\n```yaml\npaths:\n  outdir: ${paths.data_path}/pretrained_models/${experiment_name}\n  resume_checkpointdir: null \n  pretrained_checkpointdir: null\n```\n\n### 3. Edit resource requirements\n```yaml\nclusters:\n  batch_size: 2\n  worker_nodes: 1\n  gpus_per_worker: 1\n  cpus_per_gpu: 4\n  mem_per_cpu: 16000\n```\n\n### 4. Run local training job\n```bash\npython manager.py --config-name=configs/test_pretrain_4d_mae_local.yaml\n```\n\n### 5. Launch multiple training jobs or Ray Tune jobs\nTo launch multiple training jobs, set `run_type` to `multi_run` and define a `runs` list. For Ray Tune hyperparameter sweeps, set `run_type` to `tune`.\n\n## Cluster setup\n\nRunning a job on a cluster is very similar to the local setup. Override the `defaults` in your config file to match your cluster:\n\n### SLURM Setup\n```yaml\ndefaults:\n  - clusters: abc_a100\n  - paths: abc\n```\n\n### LSF Setup\n```yaml\ndefaults:\n  - clusters: janelia_h100\n  - paths: janelia\n```\n\n# Architecture Overview\n\n```\n┌─────────────────────────────────────────────────────────────────────────────┐\n│                              Ray Cluster                                    │\n│  ┌─────────────────────────────────────────────────────────────────────────┐│\n│  │                         manager.py                                      ││\n│  │                    (SLURM / LSF / Local)                                ││\n│  └─────────────────────────────────────────────────────────────────────────┘│\n│                                    │                                        │\n│           ┌────────────────────────┼────────────────────────┐               │\n│           ▼                        ▼                        ▼               │\n│  ┌─────────────────┐    ┌─────────────────┐    ┌─────────────────┐          │\n│  │ EpochBased      │    │ ParallelEpoch   │    │   Inferencer    │          │\n│  │ Trainer         │    │ BasedTrainer    │    │                 │          │\n│  │ (DeepSpeed)     │    │ (TorchTitan)    │    │ (Distributed)   │          │\n│  └─────────────────┘    └─────────────────┘    └─────────────────┘          │\n│           │                        │                        │               │\n│           └────────────────────────┼────────────────────────┘               │\n│                                    ▼                                        │\n│  ┌─────────────────────────────────────────────────────────────────────────┐│\n│  │                      Ray Data Pipeline                                  ││\n│  │  LoaderActor → SharedMemory → CollatorActor → DeviceBuffer → GPU        ││\n│  └─────────────────────────────────────────────────────────────────────────┘│\n└─────────────────────────────────────────────────────────────────────────────┘\n```\n\n# Data Pipeline\n\nThe data pipeline is built on **Ray Data** with a custom queue-based system for high-throughput data loading:\n\n```\nDatabase (Supabase/Local CSV) -\u003e LoaderActor -\u003e Host SharedMemory Buffer -\u003e CollatorActor -\u003e Device Buffer -\u003e Preprocessor -\u003e Model\n```\n\n## Databases\n\n**`ParentDatabase`** / **`SupabaseDatabase`**: Flexible database classes supporting both remote (Supabase/PostgreSQL) and local (cached CSV) data sources.\n\n- **Remote mode**: Queries Supabase via ConnectorX (Arrow path) for high-performance data fetching\n- **Local mode**: Loads from cached CSV + JSON config for offline/fast iteration\n- **Data filtering**: ROI/tile/HPF selection, occupancy thresholds, CDF thresholds, annotation filtering\n\n```python\n# Example database config\ndatasets:\n  databases:\n    _target_: data.databases.supabase_database.SupabaseDatabase\n    use_cached_hypercubes_dataframe: true\n    hypercubes_dataframe_path: ${paths.data_path}/databases/hypercubes.csv\n    server_folder_path: /groups/betzig/data\n```\n\n## Ray Dataloader\n\nThe Ray-based dataloader uses a multi-actor architecture for maximum throughput:\n\n**`LoaderActor`**: Reads hypercubes from Zarr via TensorStore into pinned shared memory buffers.\n\n**`CollatorActor`** / **`FinetuneCollatorActor`**: Transfers batches from host shared memory to GPU device buffers with optional transforms, mask extraction, and target building.\n\n## Preprocessors\n\nPreprocessors provide a unified interface for task-specific data preparation:\n\n| Preprocessor | Task | Description |\n|-------------|------|-------------|\n| `RayPreprocessor` | Pretraining | dtype normalization, masking, transforms |\n| `MultiSequenceRayPreprocessor` | DINO-style | Multiple dataset streams (global/local crops), per-stream transforms and masking |\n| `ChannelSplitPreprocessor` | Channel Split | Predicts per-channel from averaged input |\n| `UpsamplePreprocessor` | Super-resolution | NA-mask downsampling for space/time upsampling |\n| `InstanceSegmentationPreprocessor` | Instance Seg | Mask/bbox extraction, target building |\n| `SAM2VideoPreprocessor` | Video segmentation | Prompt-based video segmentation/tracking |\n\n## Patch Embeddings \u0026 Channel Encoding\n\n**PatchEmbedding** convert multi-channel volumes into token sequences. Supported input layouts include `TZYXC`, `ZYXC`, `TYXC`, and `YXC`. **ChannelAdaptivePatchEmbedding** handles variable-channel input with optional learned channel embeddings and fusion via concat or pooling.\n\n## Transforms\n\nTransforms can be applied in either the **Collator** (CPU, during data loading) or the **Preprocessor** (GPU, before model forward). Available transforms in `data/transforms/`:\n\n- **`Resize`**: Spatial resizing with mask/bbox scaling\n- **`Crop`**: Random and center cropping\n- **`Normalize`**: Percentile-based normalization\n- **`ProbabilisticChoice`**: Randomly select between transform pipelines\n- **`MultiCrop3D`**: Multi-crop augmentation for DINO-style pretraining (global + local crops)\n- **`ChannelDropout`**: Randomly drop channels by selecting a subset (configurable keep ratio)\n\n## MaskGenerator\n\nGenerates patch-level masks for self-supervised learning with explicit time/space awareness:\n\n- `BLOCKED` / `BLOCKED_TIME_ONLY` / `BLOCKED_SPACE_ONLY`: Block-based masking\n- `RANDOM` / `RANDOM_SPACE_ONLY`: MAE-style random masking\n- `BLOCKED_PATTERNED`: Deterministic time downsampling patterns\n- `BLOCKED_WITH_RANDOM_FILL`: Block-based masking with random infill\n- `DINO_IBOT`: iBOT-style patch masking for DINO pretraining\n- `HIERA_MU` / `HIERA_MU_BLOCKED`: Mask-unit-aligned masking for Hiera backbone (random or block-based)\n\n# Models\n\n## Pretraining Models\n\n| Model | Location | Description |\n|-------|----------|-------------|\n| **MAE** | `models/meta_arch/maskedautoencoder.py` | Masked Autoencoder for 3D/4D volumes (ViT or Hiera backbone) |\n| **JEPA** | `models/meta_arch/jepa.py` | Joint-Embedding Predictive Architecture, single- or multistage (ViT or Hiera backbone) |\n| **DINOv3** | `models/meta_arch/dino.py` | Self-distillation with DINO + iBOT objectives (**work in progress**) |\n\n## Detection \u0026 Segmentation\n\n| Model | Location | Description |\n|-------|----------|-------------|\n| **plainDETR** | `models/meta_arch/plainDETR.py` | 3D object detection |\n| **MaskDINO** | `models/meta_arch/maskdino.py` | 3D instance segmentation |\n| **Mask2Former** | `models/meta_arch/mask2former.py` | 3D semantic/panoptic segmentation |\n| **SAM2** | `models/meta_arch/sam.py` | Segment Anything 2 for 3D volumes with automatic mask generation and iterative point-prompt refinement |\n\n## Backbones\n\n- **ViT** (`models/backbones/vit.py`): Vision Transformer with RoPE/sincos positional encoding\n- **Hiera** (`models/backbones/hiera.py`): Hierarchical Vision Transformer with mask-unit attention and multistage pooling\n- **ConvNeXt** (`models/backbones/convnext.py`): ConvNeXt backbone\n- **DinoEncoder** (`models/backbones/dino_encoder.py`): ViT backbone for DINOv3\n- **MaskedEncoder** (`models/backbones/maskedencoder.py`): ViT-based encoder with masking support for MAE/JEPA\n- **MaskedHieraEncoder** (`models/backbones/masked_hiera_encoder.py`): Hiera-based encoder with mask-unit masking for MAE/JEPA multistage pretraining\n\n## Layers\n\nKey layer implementations:\n- **Attention** (`models/layers/attention.py`): Multi-head self-attention with flash attention, RoPE attention, mask-unit attention (Hiera), deformable attention, and memory attention (SAM2)\n- **Transformer** (`models/layers/transformer.py`): Standard and deformable transformer blocks with sequence packing support\n- **Patch Embeddings** (`models/layers/patch_embeddings.py`): 3D/4D patch embedding with multiple layout support\n- **Positional Encoding** (`models/layers/positional_encoding.py`): Sinusoidal, learned, and RoPE encodings\n- **Memory Encoders** (`models/layers/memory_encoders.py`): Memory encoder/fuser for SAM2 video tracking\n- **Matchers** (`models/layers/matcher.py`): Hungarian matcher for detection/segmentation\n\n## Kernels\n\nMAE, JEPA, MaskDINO, and PlainDETR support deformable attention.\n\n# Training\n\n## EpochBasedTrainer (DeepSpeed)\n\nStandard training loop with DeepSpeed ZeRO optimization:\n\n\n\n**Features:**\n- DeepSpeed ZeRO stages 1-3\n- Mixed precision (bf16/fp16)\n- Checkpoint saving/resuming\n- Hook-based extensibility\n\n## ParallelEpochBasedTrainer (TorchTitan)\n\nAdvanced parallelism support via TorchTitan integration (**work in progress**):\n\n\n\n**Features:**\n- **TP** (Tensor Parallelism): Shards model tensor\n- **CP** (Context Parallelism): Shards sequence dimension\n- **FSDP** (Model and Data Parallelism): FSDP-based sharding\n- **Torch Compile** support\n- **Activation checkpointing**\n\n## Hooks\n\nThe training loop is extensible via a priority-based hook system. Hooks can execute at various points: `before_train`, `before_epoch`, `before_step`, `after_backward`, `after_step`, `after_epoch`, `after_train`, and validation/test equivalents.\n\n| Hook | Purpose |\n|------|---------|\n| `IterationTimer` | Tracks step/epoch/validation timing, logs ETA |\n| `LRScheduler` | Executes LR scheduler steps, logs learning rate |\n| `WeightDecayScheduleHook` | Updates weight decay on schedule |\n| `PeriodicWriter` | Writes metrics to loggers at epoch end |\n| `PeriodicCheckpointer` | Saves checkpoints at configurable intervals |\n| `BestCheckpointer` | Reports best checkpoint to Ray for model selection |\n| `BestMetricSaver` | Tracks and updates best validation metric |\n| `TorchMemoryStats` | Logs CUDA memory usage (reserved/allocated) |\n| `TorchProfiler` | PyTorch profiler with TensorBoard traces and memory snapshots |\n| `NsysProfilerHook` | NVIDIA Nsight Systems profiling for GPU timeline analysis |\n| `EarlyStopHook` | Stops training if validation metric plateaus |\n| `EMASchedulerHook` | Updates EMA beta for target networks (JEPA) |\n| `FreeDeviceBufferHook` | Releases Ray device buffers to prevent deadlocks |\n| `AnomalyDetector` | Detects NaN/Inf losses with `torch.autograd.detect_anomaly` |\n| `GarbageCollectionHook` | Periodic synchronized GC to prevent memory fragmentation |\n| `MemoryDebugHook` | Detailed memory dumps (proc, CUDA, Ray, /dev/shm) |\n| `AdjustTimeoutHook` | Adjusts distributed timeout for long-running ops |\n\n# Logging \u0026 Experiment Tracking\n\n## Event Recording\n\nThe `EventRecorder` is the central hub for collecting metrics during training. It supports step-scoped and epoch-scoped scalars with configurable reduction methods (mean, median, sum, min, max).\n\n```python\n# Record a scalar at step scope\ntrainer.event_recorder.put_scalar(\"loss\", loss_value, scope=\"step\")\n\n# Record multiple scalars with a prefix\ntrainer.event_recorder.put_scalars(scope=\"epoch\", prefix=\"val_\", accuracy=0.95, f1=0.92)\n```\n\n## W\u0026B Integration\n\nThe `WandBEventWriter` provides seamless Weights \u0026 Biases integration.\n\n\n\n**Features:**\n- Automatic login via `WANDB_API_KEY` from `.env`\n- Custom metric namespacing (`step/*`, `epoch/*`)\n- Tags and notes for experiment organization\n\nThe `LocalEventWriter` saves metrics to CSV files for offline analysis.\n\n\n\n## Metrics Processing\n\nFor advanced training loops (TorchTitan), the `MetricsProcessor` computes detailed performance metrics (**work in progress**):\n\n- **Throughput**: Tokens per second (TPS)\n- **MFU**: Model FLOPS Utilization percentage\n- **TFLOPS**: Achieved TFLOPS\n- **Timing**: Forward/backward/optimizer step times (ms)\n- **Memory**: Peak active/reserved GPU memory, allocation retries, OOMs\n\n# Inference\n\nThe `InferencerWorker` provides distributed inference with two modes:\n\n**`stitch_volume`**: Reconstructs full volumes from hypercube predictions via all-to-all communication.\n\n**`save_local`**: Saves per-sample predictions locally (for detection/segmentation tasks).\n\nSupported tasks:\n- `detection` (plainDETR)\n- `instance_segmentation` (MaskDINO, SAM2)\n- `semantic_segmentation` (Mask2Former)\n- `dense_prediction` (upsampling, channel split)\n- `pretrain` (reconstruction)\n- `feature_extractor` (feature visualization)\n\n**SAM2 automatic mask generation** (`inference/amg.py`): Grid-based point prompting with iterative refinement, NMS, stability scoring, and crop-based tiling for full 3D volumes.\n\n# Profiling\n\nMultiple profiling tools are supported:\n\n| Profiler | Use Case |\n|----------|----------|\n| **pprof** (gperftools) | CPU profiling with `@pprof_func` / `@pprof_class` decorators |\n| **NVIDIA Nsys** | GPU timeline profiling via `NsysProfilerHook` |\n| **PyTorch Profiler** | Operator-level profiling via `TorchProfiler` hook |\n| **Memory Profiler** | PyTorch CUDA memory snapshots and `TorchMemoryStats` hook |\n| **Ray Profiler** | Ray actor/task profiling via `MemoryDebugHook` |\n\n\n\n# Configuration layout\n\nHere's what each configuration subdirectory handles:\n\n- **[`configs/models/`](configs/models)** - Model architectures (MAE, JEPA, DINO, SAM2, plainDETR, MaskDINO, Mask2Former, backbones, heads)\n- **[`configs/datasets/`](configs/datasets)** - Dataset classes, databases, and preprocessor parameters\n- **[`configs/tasks/`](configs/tasks)** - Task-specific configs (channel_split, instance_segmentation, upsample_*)\n- **[`configs/optimizers/`](configs/optimizers)** - Optimizer configurations (AdamW, LAMB, Lion, Muon)\n- **[`configs/schedulers/`](configs/schedulers)** - Learning rate and weight decay schedulers\n- **[`configs/optimizations/`](configs/optimizations)** - Model optimizations (torch.compile, activation checkpointing)\n- **[`configs/hooks/`](configs/hooks)** - Training hooks configuration\n- **[`configs/checkpoint/`](configs/checkpoint)** - Checkpointing configurations\n- **[`configs/deepspeed/`](configs/deepspeed)** - DeepSpeed ZeRO configurations\n- **[`configs/parallelism/`](configs/parallelism)** - TorchTitan parallelism settings (TP, CP, PP, DP)\n- **[`configs/clusters/`](configs/clusters)** - Cluster configurations (SLURM, LSF, local)\n- **[`configs/paths/`](configs/paths)** - Path configurations for different environments (ABC, Janelia, CoreWeave)\n- **[`configs/loggers/`](configs/loggers)** - Logging configurations (WandB, Local CSV)\n- **[`configs/profiling/`](configs/profiling)** - Profiling configurations (pprof, nsys, torch profiler)\n- **[`configs/trainer/`](configs/trainer)** - Training loop configurations\n- **[`configs/evaluation/`](configs/evaluation)** - Evaluation configurations\n- **[`configs/inference/`](configs/inference)** - Inference and prediction configurations\n- **[`configs/tune/`](configs/tune)** - Ray Tune hyperparameter sweep configurations\n- **[`configs/benchmarks/`](configs/benchmarks)** - Benchmarking configurations for throughput testing\n- **[`configs/experiments/`](configs/experiments)** - Complete experiment configs and examples\n\n# License\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at:\n\n   [Apache License 2.0](LICENSE)\n\nCopyright 2025 Cell Observatory.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcell-observatory%2Fcell_observatory_platform","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcell-observatory%2Fcell_observatory_platform","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcell-observatory%2Fcell_observatory_platform/lists"}