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!
the model trained on cifar10 which is 32*32 but test on kodim which is 768*512.
the model trained on imagenet which is resized to 128*128 but test on kodim which is 768*512.
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
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
Run(example presented in paper) on cifar10
python train.py --lr 1e-3 --epochs 1000 --batch_size 64 --channel 'AWGN' --saved ./saved --snr_list 1 4 7 13 19 --ratio_list 1/6 1/12 --dataset cifar10 --num_workers 4 --parallel True --if_scheduler True --scheduler_step_size 50
or Run(example presented in paper) on imagenet
python train.py --lr 10e-4 --epochs 300 --batch_size 32 --channel 'AWGN' --saved ./saved --dataset imagenet --num_workers 4 --parallel True
Run(example presented in paper)
python eval.py --channel 'AWGN' --saved ./saved/${mode_path} --snr 20 --test_img ${test_img_path}
- Add visualization of the model
- plot the results with different snr and ratio
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}
}