Heterogeneous Fusion of Spectrogram and Waveform with Sub-band Adaptation for Underwater Acoustic Target Recognition
SWAN fuses two branches for underwater acoustic classification:
- raw waveform branch: pre-extracted, frozen wav2vec2 last-hidden-state embeddings, time-pooled.
- spectrogram branch: processed by an sub-band adaptive spectrogram learning module and a compression module.
The two branch embeddings are projected to the same dimension and fused with a gated fusion module, followed by an MLP classifier.
SWAN/
├── swan/ # model package
│ ├── config.py # ModelConfig
│ ├── model.py # MultiModalNet
│ ├── wav2vec_branch.py
│ ├── subband_reweight.py
│ ├── mel_compression.py
│ ├── projection.py
│ ├── fusion.py # GatedFusion
│ ├── classifier.py
│ └── dataset.py
├── train.py # training script
├── evaluate.py # evaluation script
├── README.md
└── requirements.txt
pip install -r requirements.txtThe training scripts expect pre-extracted features with this layout:
<w2v_root>/
├── class_to_idx.pt
├── train/
│ └── <class>/<name>.pt # dict: {"embedding": [T, 768], "label": int}
└── test/
└── <class>/<name>.pt
<mel_root>/
├── train/
│ └── <class>/<name>.pt # dict: {"mel": [128, 157], "label": int}
└── test/
└── <class>/<name>.pt
Relative paths under train/ and test/ must match between w2v_root and mel_root.
python train.py \
--w2v_root /path/to/w2v_embeddings \
--mel_root /path/to/mel_spectrograms \
--num_classes 9 \
--dataset_name shipsear \
--epochs 100 \
--device cudaOptional flags:
--cache: preload all samples into RAM.--out_dir ./ckpts: checkpoint output directory.
Checkpoints are saved as <dataset_name>_gated_<cw|nocw>[_<run_tag>]_best.pt.
python evaluate.py \
--ckpt ./ckpts/shipsear_gated_nocw_best.pt \
--device cudaIf you use this code, please cite the corresponding work.