prashantrajbista/adjscc-cifar10

ADJSCC — Attention DeepJSCC for wireless image transmission (CIFAR-10, AWGN)

PyTorch reproduction of "Wireless Image Transmission Using Deep Source-Channel Coding With Attention Modules" (arXiv:2012.00533, Xu et al.), CIFAR-10 / AWGN, bandwidth ratio R = 1/6.

A convolutional autoencoder transmits an image directly over a noisy channel — no JPEG, no LDPC, no bit mapping. The encoder output is power-normalized and treated as complex channel symbols; AWGN is added; the decoder reconstructs the image. Attention Feature (AF) modules take the current channel SNR as an input and produce per-channel gates, so one model covers the whole 0–20 dB range instead of one model per SNR.

image 3×32×32 → encoder (5 conv + AF) → code C×8×8 → power norm → AWGN(SNR)
              → decoder (5 deconv + AF) → image 3×32×32

Code: https://github.com/prashantrajbista/wireless_image_transmission (see adjscc/).

Checkpoints

File	Model	Trained at	Params	Val PSNR @ train SNR
adjscc_r16.pt	ADJSCC (AF modules on)	SNR ~ U[0, 20] dB	11,130,956	33.13 dB @ 10 dB
base_snr1.pt	BDJSCC baseline (no AF)	fixed 1 dB	10,075,676	27.39 dB @ 1 dB
base_snr19.pt	BDJSCC baseline (no AF)	fixed 19 dB	10,075,676	36.12 dB @ 19 dB

All at C = 16 (R = k/n = C/96 = 1/6), 256 filters, AWGN. Each .pt is a dict: {"state_dict", "args", "C", "psnr"}.

Results (CIFAR-10 test set, 10,000 images, one channel draw each)

PSNR in dB, evaluated across the test SNR sweep:

Test SNR	ADJSCC	BDJSCC @1 dB	BDJSCC @19 dB
0 dB	26.52	26.65	22.34
2 dB	28.03	27.94	24.43
4 dB	29.48	28.70	26.44
6 dB	30.83	29.14	28.33
8 dB	32.06	29.42	30.08
10 dB	33.12	29.57	31.66
12 dB	34.01	29.68	33.04
14 dB	34.70	29.74	34.19
16 dB	35.22	29.78	35.12
18 dB	35.59	29.80	35.84
20 dB	35.83	29.82	36.36

The paper's main claim reproduces: the single ADJSCC model tracks the upper envelope of the fixed-SNR baselines, while each baseline collapses away from its training SNR (the 19 dB baseline loses ~4.2 dB at 0 dB; the 1 dB baseline saturates near 29.8 dB and never improves). ADJSCC stays within ~0.5 dB of the best baseline at both ends and beats both in the middle. SSIM follows the same ordering (in results/*.csv in the repo).

Plots in graphs/:

File	What
best_psnr.png	best validation PSNR per run (training curves, W&B)
train_mse.png	training MSE
test_psnr_1db.png, test_psnr_10db.png, test_psnr_19db.png	test PSNR sweeps at those channel SNRs

Usage

import torch
from huggingface_hub import hf_hub_download
from adjscc.models import DeepJSCC          # from the GitHub repo

path = hf_hub_download("prashantrajbista/adjscc-cifar10", "adjscc_r16.pt")
ck = torch.load(path, map_location="cpu")

model = DeepJSCC(C=ck["C"], F=ck["args"]["filters"],
                 attention=ck["args"]["attention"], channel="awgn")
model.load_state_dict(ck["state_dict"])
model.eval()

img = torch.rand(1, 3, 32, 32)              # CIFAR-10 image, pixels in [0,1]
with torch.no_grad():
    out = model(img, snr_db=10.0)           # channel noise applied inside forward()

attention=False for the two base_snr*.pt baselines. forward() runs encoder → power normalization → AWGN → decoder, so the channel is part of the graph; pass the SNR you want to transmit at.

Differences from the paper

This is a learning reproduction, not a replication. The core mechanism and the channel model follow the paper; the feature-extraction blocks are simplified. Read the deviation table before citing these numbers.

Faithful to the paper:

• Overall architecture shape: 5-layer strided conv encoder (32→16→8, C channels out at 8×8), 5-layer transposed-conv decoder, sigmoid output.
• AF module concept (the paper's contribution): global avg-pool → concat SNR → 2 FC layers → sigmoid gate in [0,1] → channel-wise rescale, conditioning the whole net on the channel SNR.
• Power constraint (1/k)E(zz*) ≤ 1 per complex symbol; AWGN with σ² = 10^(-SNR/10) split σ/√2 across real and imaginary parts; channel inside forward(), gradients flow through it.
• Bandwidth ratio R = k/n = C/96 for 32×32 images (k = 32C complex symbols, n = 3072).
• Training regime: MSE loss, Adam, lr 1e-4, SNR ~ U[0,20] dB resampled per batch for ADJSCC, fixed SNR for the BDJSCC baselines.
• Evaluation: integer test SNR 0–20 dB for every model.

Where this run deviates

Item	Paper	Here	Why / effect
GDN / IGDN	feature-learning module = conv + GDN (IGDN in decoder) + PReLU	not implemented — plain conv + PReLU	GDN is a significant part of the paper's absolute PSNR. Expect a gap in absolute numbers.
AF bottleneck	FC(C+1 → C/16) + ReLU → FC(C/16 → C) + sigmoid	FC(C+1 → C) → FC(C → C), no reduction	Attention overhead here is ~1.06M params instead of ~68k.
Kernel sizes	first encoder / last decoder conv 9×9	5×5 throughout	Smaller receptive field at the image boundary layers.
AF placement	after every FL module except the last of each → 8 modules	after all 5 encoder convs + 4 decoder convs → 9 modules	Minor.
Parameter count	10,690,351 (BDJSCC) / 10,758,191 (ADJSCC)	10,075,676 / 11,130,956	Follows from the three rows above — same magnitude, not the same network.
Training length	1280 epochs, batch 128	200 epochs (ADJSCC), 100 (baselines), batch 64	Compute budget. Absolute PSNR below the paper's curves; the relative ADJSCC-vs-BDJSCC behaviour, which is the paper's claim, reproduces clearly.
Augmentation	none	random horizontal flip	Minor.
Eval averaging	each test image transmitted 10×, averaged	single pass	10k test images keeps the mean stable to well under 0.1 dB, but it is not the paper's protocol.
Baseline SNRs	full set of fixed-SNR baselines	1 dB and 19 dB only	Enough to bracket the envelope and show the mismatch penalty.
Slow-fading / burst channels	reference code has slow_fading, slow_fading_eq, burst	AWGN only (channel.rayleigh raises NotImplementedError)	The paper's CIFAR-10 figures are AWGN only.
Traditional separation baseline	SPIHT + RCPC ("TJSCC"), on Kodak	not run	No classical-coding comparison here.
ImageNet / Kodak	Section IV-C versatility experiments	CIFAR-10 only	Out of scope for this reproduction.
Storage/complexity table	Table I: BDJSCC-1/2/5/10 ensembles vs ADJSCC	not run (param counts do match)	The ensemble-strategy comparison is not reproduced.
R = 1/12 curve	Fig. 6(a)	R = 1/6 only	Code supports --ratio 0.0833 (C = 8); just not trained.

Additions beyond the paper

• SSIM reported alongside PSNR. CIFAR-10 at 32×32 is too small for MS-SSIM's 5-scale pyramid, so this is single-scale SSIM — an extra metric, not a substitute.
• PSNR averaged per image (the paper's stated definition). The reference code instead averages MSE over the whole test set and converts once; the difference is small.

Citation

@article{xu2021adjscc,
  title   = {Wireless Image Transmission Using Deep Source Channel Coding With Attention Modules},
  author  = {Xu, Jialong and Ai, Bo and Chen, Wei and Yang, Ang and Sun, Peng and Rodrigues, Miguel},
  journal = {IEEE Transactions on Circuits and Systems for Video Technology},
  year    = {2022}
}

These weights are a third-party reproduction, not the authors' released models.