Deep JSCC

This implements training of Deep JSCC models for wireless image transmission as described in the paper Deep Joint Source-Channel Coding for Wireless Image Transmission by Pytorch. And there has been a Tensorflow and keras implementations .

This is my first time to use PyTorch and git to reproduce a paper, so there may be some mistakes. If you find any, please let me know. Thanks!

Update-2024.06.04

modify the train.py to omit most of the args in command line, you can just us python train.py --dataset ${dataset_name} to train the model.
add tensorboard to record the results in exp.
add the visualization/ file to visualize the result.
add bash file to run the code in parallel.

Architecture

Demo

The model trained on cifar10 which is 32*32 but test on kodim which is 768*512 (left); and the model trained on imagenet which is resized to 128*128 but test on kodim which is 768*512 (right).

Installation

conda or other virtual environment is recommended.

git clone https://github.com/chunbaobao/Deep-JSCC-PyTorch.git
cd ./Deep-JSCC-PyTorch
pip install requirements.txt

Usage

Prepare Dataset

The cifar10 dataset can be downloaded automatically by torchvision. But the imagenet dataset should be downloaded manually from ImageNet website and put in the right place, refer to dataset.py. And run:

python dataset.py

Training

The training command used to be very long, but now you can just use python train.py --dataset ${dataset_name} --channel ${channel} to train the model. The default dataset is cifar10. The parmeters can be modified in the train.py file. The default parameters are similar to the paper.

Parameters	CIFAR-10	ImageNet
`batch_size`	64	32
`init_lr`	1e-3	1e-4
`weight_decay`	5e-4	5e-4
`snr_list`	[19, 13, 7, 4, 1]	[19, 13, 7, 4, 1]
`ratio_list`	[1/6, 1/12]	[1/6, 1/12]
`if_scheduler`	True	False
`step_size`	640	N/A
`gamma`	0.1	0.1

Evaluation

The eval.py provides the evaluation of the trained model.

You may need modify slightly to evaluate the model for different snr_list and channel type in main function.

python eval.py

All training and evaluation results are saved in the ./out directory by default. The ./out directory may contain the structure as follows:

./out
├── checkpoint # trained models
│   ├── $DATASETNAME_$INNERCHANNEL_$SNR_$RATIO_$CHANNEL_TYPE_$TIMES_on_$DATE
│       ├── epoch_$num.pth
│       ├── ...
│   ├── CIFAR10_10_1.0_0.08_AWGN_13h21m37s_on_Jun_02_2024
│   ├── CIFAR10_20_7.0_0.17_Rayleigh_14h03m19s_on_Jun_03_2024
│   ├── ...
├── configs # training configurations
│   ├── $DATASETNAME_$INNERCHANNEL_$SNR_$RATIO_$CHANNEL_TYPE_$TIMES_on_$DATE
│   ├── $CIFAR10_10_4.0_0.08_AWGN_13h21m38s_on_Jun_02_2024.yaml
│   ├── ...
├── logs # training logs
│   ├── $DATASETNAME_$INNERCHANNEL_$SNR_$RATIO_$CHANNEL_TYPE_$TIMES_on_$DATE
│       ├── tensorboard logs
│   ├── ...
├── eval # evaluation results
│   ├── $DATASETNAME_$INNERCHANNEL_$SNR_$RATIO_$CHANNEL_TYPE_$TIMES_on_$DATE
│       ├── tensorboard logs
│   ├── ...

Visualization

The ./visualization directory contains the scripts for visualization of the training and evaluation results.

single_visualization.ipynb is used to get demo of the model on single image like the demo above.
plot_visualization.ipynb is used to get visualizations of the perfprmance of the model on different snr and ratio.

Results

Model results and logs for the CIFAR-10 dataset, tested under various SNR, ratio, and channel types, are available in the ./out directory. The models' performance is approximately 5dB worse than reported in the paper, likely due to the implementation reflecting a real communication system. However, the performance trends are consistent with those in the paper.

TO-DO

~~Add visualization of the model~~
~~plot the results with different snr and ratio~~
~~add Rayleigh channel~~
train on imagenet
Convert the real communication system to a complex communication system

Citation

If you find (part of) this code useful for your research, please consider citing

@misc{chunhang_Deep-JSCC,
  author = {chunhang},
  title = {a pytorch implementation of Deep JSCC},
  url ={https://github.com/chunbaobao/Deep-JSCC-PyTorch},
  year = {2023}
}

6.0 KiB Raw Blame History