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

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:

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/:

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

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.