How asunatan/hsqformer works
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: __init__.py Model entry · engine.py Engine · helpers.py Shared utils · train.py train.py
Components 19
Training pipeline
Entry point, data loading, orchestration, optimization, and evaluation.
- train.py · Main experiment entrypoint that assembles and launches training.
train.py - Data loader · Main dataset boundary for reading images and split definitions.
dataset/data_loading.py - Dataset files · File-based ultrasound classification data and split CSV used for training and evaluation.
dataset/test_3 - Augmentation · Image preprocessing and augmentation stage before model ingestion.
utils/augmentation.py - Engine · Core training and evaluation control flow.
utils/engine.py - Train/eval · Reusable training and validation loop implementations.
utils/train_eval.py - Optimizers · Optimizer and scheduler construction abstraction.
utils/optim_factory.py - Losses · Loss and criterion definitions used during training.
utils/loss.py - Metrics · Evaluation metrics and scoring utilities.
utils/metrics.py - Distributed run · Multi-GPU and submitit-based job launching support.
utils/run_with_submitit.py
Model zoo
Interchangeable backbone families and shared model infrastructure.
- Model entry · Central model import and selection boundary for experiments.
models/__init__.py - Shared utils · Common blocks and architecture-specific helper utilities shared across models.
models/helpers.py - Swin family · Swin-style classification and segmentation variants.
models/swin_transformer.py - CSWin family · CSWin-based variants including segmentation and contrastive experiments.
models/cswin.py - ConvNeXt family · ConvNeXt classification and segmentation variants.
models/convnext.py - Hiera family · Hiera classification and segmentation variants.
models/hiera.py - Liver models · Liver-focused classification and ablation architectures built for the HCC screening task.
models/LivNet.py - Other backbones · 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 · Compiled C++/CUDA window-processing acceleration used by Swin-like models.
models/kernels/window_process/swin_window_process.cpp
Connections 24
- 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.