Compression Framework for Light 3D Scene Graph Generation via Pruning-as-Search and Distillation

IEEE Transactions on Multimedia, 2026. Paper · BibTeX

*Hojun Song¹, *Chae-yeong Song¹, Dong-hun Lee¹, Heejung Choi¹, Jinwoo Jeong², Sungjei Kim³, and †Sang-hyo Park¹

• Equal contribution
  † Corresponding author

¹Kyungpook National University (KNU)
²Korea Electronics Technology Institute (KETI)
³Korea University of Technology and Education (KOREATECH)

A lightweight compression framework for GNN-based 3D scene graph generation that integrates pruning-as-search and knowledge distillation.

Dependencies

Create the conda environment directly from this repository:

scripts/setup_env.sh
conda activate 3DSGG-comp

The environment is defined in environment.yml.

Prepare the Data

A. Download 3RScan and 3DSSG-Subset annotation.

Follow the official 3DSSG preparation. The helper script can also prepare the public 3DSSG-Subset annotation:

scripts/prepare_data.sh

B. Place or link the data.

data/3RScan
data/3DSSG_subset

Expected layout:

data
  3DSSG_subset
    relationships_train.json
    relationships_validation.json
    classes.txt
    relationships.txt
    relations.txt
    train_scans.txt
    validation_scans.txt

  3RScan
    <scan_id>
      labels.instances.align.annotated.v2.ply
      semseg.v2.json
      multi_view
      ...

Train

Run the paper-level training flow:

# Environment/data check
./run.sh check --exp sgfn_baseline_repro

# Baseline teacher training before the paper compression stages
./run.sh baseline --exp sgfn_baseline_repro --device cuda

# Paper Stage 1: CO/SIE search, then Alpha model training
./run.sh stage1 --exp sgfn_baseline_repro --device cuda

# Paper Stage 2: CO2/MIE pruning, then DDR/KD fine-tuning
./run.sh stage2 --exp sgfn_baseline_repro --device cuda

./run.sh all --exp sgfn_baseline_repro --device cuda runs the same training sequence end-to-end. For long runs, use tmux or screen.

Use custom paths or configs when needed:

./run.sh all \
  --exp my_experiment \
  --device cuda \
  --config configs/sgfn.yaml \
  --scan-root /path/to/3RScan \
  --annotation-root /path/to/3DSSG_subset

run.sh uses the currently active Python environment by default. Set CONDA_ENV_NAME=<env> only when you want the wrapper to call conda run -n <env>.

Command	Purpose
check	Validate environment, data paths, and SGFN input/output wiring
baseline	Train/evaluate the baseline SGFN teacher before compression
stage1	Paper Stage 1: CO/SIE search and Alpha model training
stage2	Paper Stage 2: CO2/MIE pruning and DDR/KD fine-tuning
stage2-direct	Ablation that skips paper Stage 1 and applies Stage 2 to baseline

Supported Models

The current public run path is validated end-to-end for SGFN.

The legacy research workspace contains experiments and target definitions for other backbones used in the paper, such as SGFN-Attn, SGPN, SGGpoint, VL-SAT, and IMP. They are not exposed as complete public run.sh workflows yet. Each additional model needs its own executable adapter for baseline training, Stage 1 search/training, Stage 2 structured pruning, and DDR/KD fine-tuning.

Evaluate

Evaluate checkpoints produced by each stage:

# Baseline teacher
./run.sh evaluate-baseline --exp sgfn_baseline_repro --device cuda

# Alpha model after paper Stage 1
./run.sh evaluate-alpha --exp sgfn_baseline_repro --device cuda

# Initial pruned model before DDR/KD fine-tuning
./run.sh evaluate-pruned --exp sgfn_baseline_repro --device cuda

# Final DDR/KD model
./run.sh evaluate-final --exp sgfn_baseline_repro --device cuda

Default checkpoint paths are:

Command	Default checkpoint
evaluate-baseline	runs/<exp>/stage1/baseline_best.pt
evaluate-alpha	runs/<exp>/stage3/alpha_best.pt
evaluate-pruned	runs/<exp>/stage4/pruned_initial.pt
evaluate-final	runs/<exp>/stage5/kd_best.pt

Pass --checkpoint /path/to/checkpoint.pt to evaluate a specific checkpoint.

For internal stage mapping, output files, and the direct Stage 2 ablation details, see docs/reproduction_details.md.

Expected Runtime

Approximate runtime on an RTX 4090 with the SGFN configuration and 548 validation scenes:

Command	Main cost	Expected time
./run.sh check	Environment and data validation	1-3 min
./run.sh baseline	100-epoch baseline training plus validation every 20 epochs	10-12 h
./run.sh stage1	CO/SIE evaluation passes plus Alpha training	16-19 h
./run.sh stage2	CO2/MIE evaluation passes plus DDR/KD fine-tuning	16-19 h
./run.sh stage2-direct	CO2/MIE evaluation passes plus DDR/KD fine-tuning, without Alpha training	11-13 h
./run.sh evaluate-*	Full validation evaluation	30-35 min

Stage 1 and Stage 2 are long because CO/SIE/CO2/MIE require repeated validation passes. See docs/reproduction_details.md for the internal command mapping.

Citation

If you find this work useful, please cite:

@ARTICLE{11329462,
  author={Song, Hojun and Song, Chae-yeong and Lee, Dong-hun and Choi, Heejung and Jeong, Jinwoo and Kim, Sungjei and Park, Sang-hyo},
  journal={IEEE Transactions on Multimedia},
  title={Compression Framework for Light 3D Scene Graph Generation via Pruning-As-Search and Distillation},
  year={2026},
  volume={},
  number={},
  pages={1-13},
  keywords={Three-dimensional displays;Computational modeling;Accuracy;Point cloud compression;Semantics;Model compression;Electronic mail;Adaptation models;Knowledge transfer;Image coding;3D scene graph generation;graph neural networks;model compression;pruning;knowledge distillation},
  doi={10.1109/TMM.2026.3651094}}

Acknowledgement

This work was supported in part by the Institute of Information and Communications Technology Planning and Evaluation (IITP) funded by the Korean government (MSIT), in part by the BK21 FOUR project funded by the Ministry of Education, and in part by the National Research Foundation of Korea (NRF).

Our work builds on the following repository: VL-SAT.