Skip to content

Repository files navigation

📡 CSI-CLIP: A CIR-CSI Contrastive Channel Foundation Model

arXiv IEEE License: CC BY-NC 4.0 PyTorch CFM

This is the official code release for CSI-CLIP, the contrastive channel foundation model introduced in A MIMO Wireless Channel Foundation Model via CIR-CSI Consistency (arXiv:2502.11965). CSI-CLIP learns transferable wireless-channel representations by aligning channel frequency response (CFR/CSI) and channel impulse response (CIR) views with CLIP-style contrastive learning.

✨ Contributions

CSI-CLIP is designed as a reusable pre-training framework for wireless channels:

  • CIR-CSI consistency: aligns frequency-domain CFR/CSI features and time-domain CIR features derived from the same channel.
  • Contrastive channel modeling: treats naturally paired CIR and CSI samples as positive pairs and other in-batch pairs as negatives.
  • Scenario-level pre-training: supports multi-scenario training with held-out validation scenarios.
  • Open release templates: provides path-anonymized scripts for local data, checkpoints, and experiment outputs.

🧩 Method Overview

CSI-CLIP framework

The release pre-trains two ResNet-50 encoders, one for CFR/CSI and one for CIR. CFR samples are converted to two-channel real/imaginary tensors, transformed to CIR through IFFT, and optimized with a symmetric CLIP-style contrastive loss between the paired modalities.

📁 Repository Contents

  • pretrain.py: distributed CSI-CLIP contrastive pre-training.
  • downstream/train.py: downstream training for positioning, beam management, and LOS/NLOS classification.
  • downstream/eval.py: downstream checkpoint evaluation.
  • pretrain.sh: path-anonymized run template.
  • positioning.sh, beam_management.sh, los.sh: downstream run templates.
  • dataset.py: scenario dataset loader for cfr.npy.
  • downstream/datasets.py: DeepMIMO-style downstream dataset loaders.
  • model.py: ResNet-50 encoders and projection heads.
  • downstream/models.py: downstream heads.
  • loss.py: CLIP-style symmetric contrastive loss.
  • augmentations.py: CFR/CIR preprocessing transforms without extra augmentation.
  • utils.py: distributed training utilities.

Data files, checkpoints, TensorBoard logs, and experiment outputs are not included in this release.

⚙️ Environment

Required packages include:

  • Python 3.8+
  • PyTorch
  • torchvision
  • numpy
  • tqdm

Example installation:

pip install torch torchvision numpy tqdm

🗂️ Data Layout

Prepare data locally and pass paths through environment variables or command-line arguments.

Pre-training data:

DATA_ROOT/
  scenario_a/
    cfr.npy
  scenario_b/
    cfr.npy

By default, dataset.py loads each cfr.npy as a complex array and reshapes it to (-1, 256, 256). It applies per-sample, per-channel min-max normalization and then builds the aligned CIR view with IFFT. Adapt dataset.py if your local CFR dimensions differ.

Downstream data:

DATA_ROOT/
  scenario_a/
    train_csi.npy
    val_csi.npy
    train_pos.npy   # positioning, shape [N, >=2]
    val_pos.npy
    train_bm.npy    # beam management class id, shape [N]
    val_bm.npy
    train_los.npy   # LOS/NLOS class id, shape [N]
    val_los.npy

Downstream loaders apply the same two-channel real/imaginary conversion and per-sample min-max normalization as pre-training.

🚀 How to Run

🧠 Pre-train CSI-CLIP

DATA_ROOT=/path/to/pretrain_data \
OUTPUT_ROOT=./outputs \
GPUS=0,1,2,3 \
NPROC_PER_NODE=4 \
bash pretrain.sh

Select validation scenarios with a comma-separated list:

DATA_ROOT=/path/to/pretrain_data \
VAL_SCENARIOS=scenario_a,scenario_b \
bash pretrain.sh

Equivalent direct launch:

torchrun \
  --nproc_per_node=4 \
  --master_addr=localhost \
  --master_port=12355 \
  pretrain.py \
  --data_root /path/to/pretrain_data \
  --val_scenarios scenario_a,scenario_b \
  --output_dir ./outputs/pretrain_r50_deepmimo

🔧 Train Downstream Tasks

Use PRETRAINED_CKPT to initialize the downstream CSI encoder from a CSI-CLIP checkpoint. Set FREEZE_ENCODER=1 to train only the downstream head.

Positioning:

DATA_ROOT=/path/to/deepmimo_data \
SCENARIO=scenario_a \
PRETRAINED_CKPT=/path/to/csi_clip/best.pth \
bash positioning.sh

Beam management:

DATA_ROOT=/path/to/deepmimo_data \
SCENARIO=scenario_a \
NUM_CLASSES=64 \
PRETRAINED_CKPT=/path/to/csi_clip/best.pth \
bash beam_management.sh

LOS/NLOS classification:

DATA_ROOT=/path/to/deepmimo_data \
SCENARIO=scenario_a \
PRETRAINED_CKPT=/path/to/csi_clip/best.pth \
bash los.sh

Equivalent direct launch:

python -m downstream.train \
  --task positioning \
  --data_root /path/to/deepmimo_data \
  --scenario scenario_a \
  --pretrained_ckpt /path/to/csi_clip/best.pth \
  --output_dir ./outputs/downstream/positioning/scenario_a

📊 Evaluate Downstream Checkpoints

python -m downstream.eval \
  --task beam \
  --data_root /path/to/deepmimo_data \
  --scenario scenario_a \
  --checkpoint ./outputs/downstream/beam/scenario_a/best.pth

📧 Contact

If you have any questions, please feel free to contact Jun Jiang at Jun.Jiang25@student.xjtlu.edu.cn.

🙏 Acknowledgement

This codebase uses PyTorch and torchvision components, and includes utility code adapted from public self-supervised learning implementations released by Meta/Facebook AI. The dual-encoder logits and symmetric cross-entropy training pattern follow the public OpenAI CLIP implementation. Original attribution notices are retained in source files where applicable.

📜 License

This project is released for research and other non-commercial use only under the Creative Commons Attribution-NonCommercial 4.0 International license. Commercial use is prohibited unless prior written authorization is obtained from the authors. See LICENSE for details.

📝 Citation

If you find this work helpful, please consider citing:

@inproceedings{jiang2025csi_clip,
  title={A MIMO Wireless Channel Foundation Model via CIR-CSI Consistency},
  author={Jiang, Jun and Yu, Wenjun and Li, Yunfan and Gao, Yuan and Xu, Shugong},
  booktitle={2025 IEEE International Conference on Machine Learning for Communication and Networking (ICMLCN)},
  pages={1--6},
  year={2025},
  doi={10.1109/ICMLCN64995.2025.11140262}
}

Some other related papers and resources:

@article{jiang2026csimae,
  title={CSI-MAE: A Masked Autoencoder-based Channel Foundation Model},
  author={Jiang, Jun and Ruan, Xiaolong and Xu, Shugong},
  journal={arXiv preprint arXiv:2601.03789},
  year={2026}
}

@article{jiang2025cfmsurvey,
  title={Towards Channel Foundation Models (CFMs): Motivations, Methodologies and Opportunities},
  author={Jiang, Jun and Gao, Yuan and Wu, Xinyi and Xu, Shugong},
  journal={arXiv preprint arXiv:2507.13637},
  year={2025}
}

About

No description, website, or topics provided.

Resources

Stars

20 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages