# asunatan/hsqformer architecture

> Architecture diagram of the GitHub repository asunatan/hsqformer, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-19). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/asunatan/hsqformer
- Repository: https://github.com/asunatan/hsqformer

## Overview

1. **What this repository is**
   - This is a research codebase for the **HSQformer** project: a medical imaging model family aimed at **ultrasound liver screening / hepatocellular carcinoma-related classification**.
   - The repository is organized as an **experiment workspace** rather than a polished product; the README explicitly says the code is not yet well organized.

2. **Primary execution path**
   - The main entrypoint is **`train.py`**.
   - Training behavior is assembled through the **`utils/`** package, which likely handles the full loop: data loading, augmentation, optimization, metrics, logging, and evaluation.

3. **Data subsystem**
   - **`dataset/data_loading.py`** is the main data boundary.
   - The visible dataset artifacts suggest a **binary classification setup** with class folders such as **`Benign/`** and **`Malignant/`**, plus a CSV file (**`test3.csv`**) for test/validation splits.
   - The repository also contains a test directory (**`dataset/test_3/`**) indicating a file-based image dataset, probably used for offline evaluation.

4. **Model subsystem: a large experiment zoo**
   - The **`models/`** directory contains many architectures, showing this repo is built for **model comparison and ablation**, not just one network.
   - Core families include:
     - **Transformer-style vision backbones**: `swin_transformer.py`, `swin_transformer_v2.py`, `cswin.py`, `pvt_v2.py`, `mpvit.py`, `hiera.py`, `focalnet.py`, `davit.py`
     - **ConvNeXt variants**: `convnext.py`, `convnextv2_segmentation.py`, `convnext_segmentaion.py`
     - **Liver-specific composites**: `LivNet.py`, `liver_convnext.py`, `liver_cross_net.py`, `liver_ablation.py`, `liver_net_conv_classification.py`
     - **Other experimental backbones**: `moganet.py`, `thynet.py`, `modeling_switch_transformers.py`
   - This suggests the architecture diagram should show a **shared training pipeline feeding multiple interchangeable model implementations**.

5. **Task-specific branches: classification and segmentation**
   - Some files are explicitly segmentation-oriented: `swin_transformer_seg.py`, `swin_transformer_segmentation.py`, `cswin_transformer_segmentation.py`, `hiera_segmentation.py`, `convnext_segmentaion.py`, `convnextv2_segmentaion.py`.
   - Others are classification-oriented, especially the liver-specific models.
   - So the repo supports at least **two major task flows**:
     - **Classification** for benign vs malignant prediction
     - **Segmentation / transformer-assisted spatial modeling** for related experiments or feature extraction

6. **Shared model infrastructure**
   - **`models/helpers.py`**, **`models/hiera_utils.py`**, and model-specific utility modules provide common blocks and glue code.
   - These files likely encapsulate reusable layers, initialization, attention/window helpers, and architecture-specific building blocks shared across variants.

7. **Custom low-level acceleration**
   - The `models/kernels/window_process/` subtree is a separate **native extension boundary**.
   - It includes **C++/CUDA code** (`swin_window_process.cpp`, `swin_window_process_kernel.cu`) plus `setup.py`, indicating custom GPU acceleration for **window-based attention processing** used by Swin-like models.
   - This is an important architectural boundary because some model variants depend on compiled extensions rather than pure PyTorch.

8. **Training loop and orchestration utilities**
   - **`utils/engine.py`**, **`utils/train_eval.py`**, and **`utils/misc.py`** likely define the core training/evaluation control flow.
   - **`utils/distributed_utils.py`** and **`utils/run_with_submitit.py`** show the code supports **distributed execution / job submission**, likely for multi-GPU or cluster training.

9. **Losses, matching, metrics, and regularization**
   - The repo separates optimization logic from model definitions:
     - **`utils/loss.py`**, `utils/criterion.py`
     - **`utils/matcher.py`** for alignment/matching logic if using detection/segmentation-style supervision
     - **`utils/metrics.py`** for evaluation scores
     - **`utils/regularizers.py`** for auxiliary training constraints
   - This makes the training pipeline modular: model output flows into reusable objective and evaluation components.

10. **Data augmentation and feature processing**
   - **`utils/augmentation.py`** handles input transformations.
   - **`utils/point_features.py`** and **`utils/functional.py`** suggest additional feature engineering or geometry-related operations for specialized experiments.
   - The architecture likely applies augmentation before feeding images into the selected backbone.

11. **Optimization layer**
   - **`utils/optim_factory.py`** indicates the optimizer/scheduler setup is abstracted away from training code.
   - This supports swapping optimizers and tuning experiments without changing the model code.

12. **Evaluation and reporting boundary**
   - The pipeline likely ends in **metrics computation** and possibly saved checkpoints/weights.
   - The README references released weights and a public dataset on Hugging Face, so the project has an external artifact boundary: **training inputs, pretrained weights, and evaluation data are managed separately from code**.

13. **External dependencies that matter architecturally**
   - **PyTorch** is the underlying runtime, given the model naming and CUDA extension layout.
   - **CUDA/C++ extension tooling** is required for window-process acceleration.
   - **Submitit**-style distributed job launching is part of the training architecture.
   - Dataset release and weights are hosted externally on **Hugging Face** and **Google Drive**, but these are distribution channels rather than core runtime dependencies.

14. **How to draw the architecture**
   - A useful diagram would show:
     - **Dataset files / CSVs** on the left
     - **`dataset/data_loading.py` + augmentation** as the input pipeline
     - A **shared training driver (`train.py` + `utils/`)**
     - A **model selection layer** feeding one of many architectures in `models/`
     - Optional **custom CUDA window kernels** used by Swin-like models
     - **Loss/metrics/optimizer** modules closing the loop
     - **Distributed submission utilities** as an outer orchestration layer
   - The central design theme is **one training pipeline, many interchangeable backbones, with custom GPU ops for transformer efficiency**.

Main entry points: [`models/__init__.py`](https://github.com/asunatan/hsqformer/blob/main/models/__init__.py) (Model entry), [`utils/engine.py`](https://github.com/asunatan/hsqformer/blob/main/utils/engine.py) (Engine), [`models/helpers.py`](https://github.com/asunatan/hsqformer/blob/main/models/helpers.py) (Shared utils), [`train.py`](https://github.com/asunatan/hsqformer/blob/main/train.py) (train.py)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_pipeline["Training pipeline"]
  node_train_py["train.py<br/>entrypoint<br/>[train.py]"]
  node_data_loading["Data loader<br/>dataset boundary<br/>[data_loading.py]"]
  node_dataset_files[("Dataset files<br/>image dataset")]
  node_augmentation["Augmentation<br/>input transforms<br/>[augmentation.py]"]
  node_engine["Engine<br/>training core<br/>[engine.py]"]
  node_train_eval["Train/eval<br/>loop logic<br/>[train_eval.py]"]
  node_optim_factory["Optimizers<br/>optimization<br/>[optim_factory.py]"]
  node_losses["Losses<br/>objective<br/>[loss.py]"]
  node_metrics["Metrics<br/>evaluation<br/>[metrics.py]"]
  node_distributed["Distributed run<br/>orchestration"]
end

subgraph group_models["Model zoo"]
  node_model_router["Model entry<br/>model registry<br/>[__init__.py]"]
  node_shared_utils["Shared utils<br/>model helpers<br/>[helpers.py]"]
  node_swin_family["Swin family<br/>vision transformer"]
  node_cswin_family["CSWin family<br/>vision transformer<br/>[cswin.py]"]
  node_convnext_family["ConvNeXt family<br/>cnn backbone<br/>[convnext.py]"]
  node_hiera_family["Hiera family<br/>vision transformer<br/>[hiera.py]"]
  node_liver_models["Liver models<br/>task-specific models<br/>[LivNet.py]"]
  node_other_backbones["Other backbones<br/>experimental models<br/>[pvt_v2.py]"]
end

subgraph group_native["Native ops"]
  node_window_kernel["Window kernel<br/>cuda extension"]
end

node_train_py -->|"loads data"| node_data_loading
node_train_py -->|"starts loop"| node_engine
node_train_py -->|"selects model"| node_model_router
node_train_py -.->|"launches jobs"| node_distributed
node_data_loading -->|"reads splits"| node_dataset_files
node_data_loading -->|"preprocesses"| node_augmentation
node_engine -->|"drives loop"| node_train_eval
node_engine -->|"builds optimizers"| node_optim_factory
node_engine -->|"applies objective"| node_losses
node_engine -->|"reports scores"| node_metrics
node_train_eval -->|"fetches batches"| node_data_loading
node_train_eval -->|"runs model"| node_model_router
node_model_router -->|"shares blocks"| node_shared_utils
node_model_router -->|"provides"| node_swin_family
node_model_router -->|"provides"| node_cswin_family
node_model_router -->|"provides"| node_convnext_family
node_model_router -->|"provides"| node_hiera_family
node_model_router -->|"provides"| node_liver_models
node_model_router -->|"provides"| node_other_backbones
node_swin_family -.->|"uses"| node_window_kernel
node_cswin_family -->|"reuses"| node_shared_utils
node_hiera_family -->|"reuses"| node_shared_utils
node_convnext_family -->|"reuses"| node_shared_utils
node_liver_models -->|"reuses"| node_shared_utils

click node_train_py "https://github.com/asunatan/hsqformer/blob/main/train.py"
click node_data_loading "https://github.com/asunatan/hsqformer/blob/main/dataset/data_loading.py"
click node_dataset_files "https://github.com/asunatan/hsqformer/tree/main/dataset/test_3"
click node_augmentation "https://github.com/asunatan/hsqformer/blob/main/utils/augmentation.py"
click node_engine "https://github.com/asunatan/hsqformer/blob/main/utils/engine.py"
click node_train_eval "https://github.com/asunatan/hsqformer/blob/main/utils/train_eval.py"
click node_optim_factory "https://github.com/asunatan/hsqformer/blob/main/utils/optim_factory.py"
click node_losses "https://github.com/asunatan/hsqformer/blob/main/utils/loss.py"
click node_metrics "https://github.com/asunatan/hsqformer/blob/main/utils/metrics.py"
click node_distributed "https://github.com/asunatan/hsqformer/blob/main/utils/run_with_submitit.py"
click node_model_router "https://github.com/asunatan/hsqformer/blob/main/models/__init__.py"
click node_shared_utils "https://github.com/asunatan/hsqformer/blob/main/models/helpers.py"
click node_swin_family "https://github.com/asunatan/hsqformer/blob/main/models/swin_transformer.py"
click node_cswin_family "https://github.com/asunatan/hsqformer/blob/main/models/cswin.py"
click node_convnext_family "https://github.com/asunatan/hsqformer/blob/main/models/convnext.py"
click node_hiera_family "https://github.com/asunatan/hsqformer/blob/main/models/hiera.py"
click node_liver_models "https://github.com/asunatan/hsqformer/blob/main/models/LivNet.py"
click node_other_backbones "https://github.com/asunatan/hsqformer/blob/main/models/pvt_v2.py"
click node_window_kernel "https://github.com/asunatan/hsqformer/blob/main/models/kernels/window_process/swin_window_process.cpp"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train_py,node_data_loading,node_dataset_files,node_augmentation,node_engine,node_train_eval,node_optim_factory,node_losses,node_metrics,node_distributed toneBlue
class node_model_router,node_shared_utils,node_swin_family,node_cswin_family,node_convnext_family,node_hiera_family,node_liver_models,node_other_backbones toneAmber
class node_window_kernel toneMint
```

## Components

### Training pipeline

Entry point, data loading, orchestration, optimization, and evaluation.

- [train.py](https://github.com/asunatan/hsqformer/blob/main/train.py): entrypoint · Main experiment entrypoint that assembles and launches training. · `train.py`
- [Data loader](https://github.com/asunatan/hsqformer/blob/main/dataset/data_loading.py): dataset boundary · Main dataset boundary for reading images and split definitions. · `dataset/data_loading.py`
- [Dataset files](https://github.com/asunatan/hsqformer/tree/main/dataset/test_3): image dataset · File-based ultrasound classification data and split CSV used for training and evaluation. · `dataset/test_3`
- [Augmentation](https://github.com/asunatan/hsqformer/blob/main/utils/augmentation.py): input transforms · Image preprocessing and augmentation stage before model ingestion. · `utils/augmentation.py`
- [Engine](https://github.com/asunatan/hsqformer/blob/main/utils/engine.py): training core · Core training and evaluation control flow. · `utils/engine.py`
- [Train/eval](https://github.com/asunatan/hsqformer/blob/main/utils/train_eval.py): loop logic · Reusable training and validation loop implementations. · `utils/train_eval.py`
- [Optimizers](https://github.com/asunatan/hsqformer/blob/main/utils/optim_factory.py): optimization · Optimizer and scheduler construction abstraction. · `utils/optim_factory.py`
- [Losses](https://github.com/asunatan/hsqformer/blob/main/utils/loss.py): objective · Loss and criterion definitions used during training. · `utils/loss.py`
- [Metrics](https://github.com/asunatan/hsqformer/blob/main/utils/metrics.py): evaluation · Evaluation metrics and scoring utilities. · `utils/metrics.py`
- [Distributed run](https://github.com/asunatan/hsqformer/blob/main/utils/run_with_submitit.py): orchestration · Multi-GPU and submitit-based job launching support. · `utils/run_with_submitit.py`

### Model zoo

Interchangeable backbone families and shared model infrastructure.

- [Model entry](https://github.com/asunatan/hsqformer/blob/main/models/__init__.py): model registry · Central model import and selection boundary for experiments. · `models/__init__.py`
- [Shared utils](https://github.com/asunatan/hsqformer/blob/main/models/helpers.py): model helpers · Common blocks and architecture-specific helper utilities shared across models. · `models/helpers.py`
- [Swin family](https://github.com/asunatan/hsqformer/blob/main/models/swin_transformer.py): vision transformer · Swin-style classification and segmentation variants. · `models/swin_transformer.py`
- [CSWin family](https://github.com/asunatan/hsqformer/blob/main/models/cswin.py): vision transformer · CSWin-based variants including segmentation and contrastive experiments. · `models/cswin.py`
- [ConvNeXt family](https://github.com/asunatan/hsqformer/blob/main/models/convnext.py): cnn backbone · ConvNeXt classification and segmentation variants. · `models/convnext.py`
- [Hiera family](https://github.com/asunatan/hsqformer/blob/main/models/hiera.py): vision transformer · Hiera classification and segmentation variants. · `models/hiera.py`
- [Liver models](https://github.com/asunatan/hsqformer/blob/main/models/LivNet.py): task-specific models · Liver-focused classification and ablation architectures built for the HCC screening task. · `models/LivNet.py`
- [Other backbones](https://github.com/asunatan/hsqformer/blob/main/models/pvt_v2.py): experimental models · Additional experimental architectures such as PVT, MPViT, FocalNet, DaViT, MoGaNet, ThyNet, and Switch Transformers. · `models/pvt_v2.py`

### Native ops

Custom compiled acceleration for window-based attention.

- [Window kernel](https://github.com/asunatan/hsqformer/blob/main/models/kernels/window_process/swin_window_process.cpp): cuda extension · Compiled C++/CUDA window-processing acceleration used by Swin-like models. · `models/kernels/window_process/swin_window_process.cpp`

## Connections

- train.py → Data loader: loads data · The main script binds the dataset boundary into the experiment.
- train.py → Engine: starts loop · The entrypoint delegates execution to the training engine.
- train.py → Model entry: selects model · Experiment configuration chooses one backbone family through the model registry.
- train.py → Distributed run: launches jobs · Cluster or multi-GPU execution can be launched through submitit utilities.
- Data loader → Dataset files: reads splits · The loader consumes image folders and CSV split artifacts.
- Data loader → Augmentation: preprocesses · Loaded samples flow through augmentation before training.
- Engine → Train/eval: drives loop · The engine coordinates the reusable train/validation loops.
- Engine → Optimizers: builds optimizers · Optimization setup is abstracted away from the core loop.
- Engine → Losses: applies objective · Training uses modular loss and criterion definitions.
- Engine → Metrics: reports scores · Validation and reporting depend on shared metric utilities.
- Train/eval → Data loader: fetches batches · The loop pulls minibatches from the dataset pipeline.
- Train/eval → Model entry: runs model · Batches are forwarded through the selected backbone.
- Model entry → Shared utils: shares blocks · Many model variants reuse common helper code.
- Model entry → Swin family: provides · Swin variants form one of the main interchangeable families.
- Model entry → CSWin family: provides · CSWin variants are another backbone family.
- Model entry → ConvNeXt family: provides · ConvNeXt variants are available for comparison experiments.
- Model entry → Hiera family: provides · Hiera variants participate in the model zoo.
- Model entry → Liver models: provides · Task-specific liver classifiers and ablations are available.
- Model entry → Other backbones: provides · Additional experimental backbones are exposed through the same interface.
- Swin family → Window kernel: uses · Swin-style models rely on compiled window-processing acceleration.
- CSWin family → Shared utils: reuses · Transformer variants share helper infrastructure.
- Hiera family → Shared utils: reuses · Hiera implementations build on common utilities.
- ConvNeXt family → Shared utils: reuses · ConvNeXt variants reuse shared model utilities where needed.
- Liver models → Shared utils: reuses · Liver-specific architectures compose shared blocks and helpers.