From 6c8471ece042c0d15dc65ea9ae16cc30ca86afa9 Mon Sep 17 00:00:00 2001 From: KiHoLee Date: Mon, 3 Aug 2026 12:49:53 +0900 Subject: [PATCH] Match figure styling and labels to the submitted manuscript Enlarge the in-canvas fonts and line weights of the result figures so that they stay legible at the printed column width, split the Fig. 5 convergence curve into a pre-meta adaptation entry and the proposed MAML entry, and rename the autoencoder legend to match the table row. Add the analytic complexity replot behind Fig. 4, which was missing from the repository, and correct the table numbering in the README. --- .gitignore | 1 + README.md | 14 +- c13_mnist.py | 1179 +++++++++++++++++++------------------- c18_mnist_doppler.py | 15 +- c19_mnist_epoch.py | 40 +- c20_bert.py | 15 +- fig_replot_complexity.py | 41 ++ 7 files changed, 697 insertions(+), 608 deletions(-) create mode 100755 fig_replot_complexity.py diff --git a/.gitignore b/.gitignore index eee60b2..797a2c9 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ data_mnist/ results_bert/bert_feats.pt __pycache__/ *.pyc +fig/ diff --git a/README.md b/README.md index 3d5277e..0a681f6 100755 --- a/README.md +++ b/README.md @@ -19,8 +19,9 @@ fixed random seeds. | `c18_mnist_doppler.py` | Doppler sweep behind Fig. 7, with task-conditional decoder-side adaptation at every operating point. | | `c19_mnist_epoch.py` | Convergence study behind Fig. 5, recording the task-adapted SER at every training checkpoint. | | `c22_maml_epoch.py` | Budget-matched meta-training trajectory that forms the right-hand segment of Fig. 5. | -| `c21_mnist_flat.py` | Flat Rayleigh study behind Table III, including the decoder-side first-order MAML variant. | +| `c21_mnist_flat.py` | Flat Rayleigh study behind Table II, including the decoder-side first-order MAML variant. | | `c20_bert.py` | Concluding BERT/AG News text study behind Fig. 8, from feature caching to training and evaluation. | +| `fig_replot_complexity.py` | Analytic complexity comparison behind Fig. 4, for the softmax and the signed realizations. | ## Reproducing the figures @@ -28,14 +29,21 @@ Each figure regenerates from the stored CSV results without retraining. ```bash +python3 fig_replot_complexity.py # Fig. 4 python3 c19_mnist_epoch.py --mode fig # Fig. 5 python3 c13_mnist.py --mode fig # Fig. 6 python3 c18_mnist_doppler.py --mode fig # Fig. 7 python3 c20_bert.py --mode fig # Fig. 8 ``` -Table III values are stored in `results_mnist/mnist_flat.csv`, and -Table IV values come from `results_mnist/mnist_results.csv` and +Fig. 4 is analytic and needs no stored results. In Fig. 5 the green +curve is split into two legend entries: left of the dotted line the +five-step adaptation starts from the jointly trained model and no +meta-training has taken place, and right of it the receiver +meta-trains within the same total step budget. + +Table II values are stored in `results_mnist/mnist_flat.csv`, and +Table III values come from `results_mnist/mnist_results.csv` and `results_mnist/mnist_doppler.csv`. ## Rerunning the experiments diff --git a/c13_mnist.py b/c13_mnist.py index 1200406..2eb8914 100755 --- a/c13_mnist.py +++ b/c13_mnist.py @@ -1,586 +1,593 @@ -#!/usr/bin/env python3 -# ============================================================ -# c13_mnist.py -# -# Real-data validation: multi-user semantic transmission of MNIST -# images over the time-varying frequency-selective channel of c11. -# -# Proposed: shared CNN semantic encoder -> d-dim embedding per user -# -> learnable user masks + over-the-air superposition -> OFDM -# frame (pilot + data, CSI aging) -> aging-aware LMMSE + RMS norm -# -> signed user-wise attention demux -> shared 10-class head. -# Metric: semantic error rate, P(recovered class != true label). -# -# Conventional scheme: transmit-side classification with the same -# CNN trunk, then digital transmission of the 4-bit class index as -# two QPSK symbols over the user's 16 comb subcarriers (8-fold -# repetition each, MRC with genie or aged pilot CSI). -# -# Usage: -# python3 c13_mnist.py --mode train -# python3 c13_mnist.py --mode train-tx-cls -# python3 c13_mnist.py --mode eval -# python3 c13_mnist.py --mode fig -# ============================================================ - -import argparse -import csv -import math -import os -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from torchvision import datasets, transforms - -import c11_doppler_csi as base - - -# ------------------------------------------------------------ -# Models -# ------------------------------------------------------------ -class CNNTrunk(nn.Module): - def __init__(self, out_dim): - super().__init__() - self.net = nn.Sequential( - nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.ReLU(inplace=True), - nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU(inplace=True), - nn.Flatten(), - nn.Linear(32 * 7 * 7, out_dim), - ) - - def forward(self, x): - return self.net(x) - - -class MnistSemanticMA(nn.Module): - """CNN encoder + masks + signed user-wise attention + classifier.""" - - def __init__(self, U, d=128, hidden=256, n_cls=10, score_hidden=64): - super().__init__() - self.U, self.d = U, d - self.encoder = CNNTrunk(d) - self.masks = nn.Parameter(torch.randn(U, d)) - self.score = nn.Sequential( - nn.Linear(U * U, score_hidden), nn.ReLU(inplace=True), - nn.Linear(score_hidden, score_hidden), nn.ReLU(inplace=True), - nn.Linear(score_hidden, U * U), - ) - self.cls = nn.Sequential( - nn.Linear(2 * d, hidden), nn.ReLU(inplace=True), - nn.Linear(hidden, n_cls), - ) - - def tx(self, imgs, params=None): - # imgs: (B,U,1,28,28) - B = imgs.shape[0] - e = self.encoder(imgs.reshape(B * self.U, 1, 28, 28)).view(B, self.U, self.d) - m = F.normalize(self.masks, dim=1) - y = (e * m.unsqueeze(0)).sum(dim=1) - y = y / torch.sqrt(torch.mean(y ** 2, dim=1, keepdim=True) + 1e-12) - return y, m - - def rx(self, Yeq, m, params=None): - B = Yeq.shape[0] - R = Yeq.unsqueeze(1) * m.unsqueeze(0) - phi = torch.cat([R.real, R.imag], dim=-1) - T = torch.bmm(phi, phi.transpose(1, 2)) / (2 * self.d) - w = self.score(T.reshape(B, -1)).view(B, self.U, self.U) - W = torch.eye(self.U, device=phi.device).unsqueeze(0) + w - z = torch.bmm(W, phi) - return self.cls(z) # (B,U,10) - - - -class MnistTransformerMA(MnistSemanticMA): - """SOTA variant: Transformer separation on the equalized features.""" - - def __init__(self, U, d=128, hidden=256, n_cls=10, n_heads=4, n_layers=2): - super().__init__(U, d, hidden, n_cls) - import torch.nn as nn - layer = nn.TransformerEncoderLayer(d_model=hidden, nhead=n_heads, - dim_feedforward=2 * hidden, - batch_first=True) - self.inp = nn.Linear(2 * d, hidden) - self.sep = nn.TransformerEncoder(layer, num_layers=n_layers) - self.out = nn.Linear(hidden, n_cls) - - def rx(self, Yeq, m, params=None): - import torch - R = Yeq.unsqueeze(1) * m.unsqueeze(0) - phi = torch.cat([R.real, R.imag], dim=-1) - return self.out(self.sep(self.inp(phi))) - - -class MnistPerUserAE(nn.Module): - """DeepMA-style per-user AE: dedicated per-user CNN encoders (no - shared masks) and dedicated per-user decoder heads, superposed on - the identical full-band resource with unit transmit power.""" - - def __init__(self, U, d=128, hidden=256, n_cls=10): - super().__init__() - self.U, self.d = U, d - self.encs = nn.ModuleList([CNNTrunk(d) for _ in range(U)]) - self.heads = nn.ModuleList([ - nn.Sequential(nn.Linear(2 * d, hidden), nn.ReLU(inplace=True), - nn.Linear(hidden, n_cls)) - for _ in range(U) - ]) - - def tx(self, imgs, params=None): - e = torch.stack([self.encs[u](imgs[:, u]) for u in range(self.U)], - dim=1) # (B,U,d) - y = e.sum(dim=1) - y = y / torch.sqrt(torch.mean(y ** 2, dim=1, keepdim=True) + 1e-12) - return y, None - - def rx(self, Yeq, m, params=None): - phi = torch.cat([Yeq.real, Yeq.imag], dim=-1) # (B,2d) - return torch.stack([h(phi) for h in self.heads], dim=1) - - -class TxClassifier(nn.Module): - """Transmit-side classifier for the conventional digital chain.""" - - def __init__(self, d=128, n_cls=10): - super().__init__() - self.trunk = CNNTrunk(d) - self.head = nn.Linear(d, n_cls) - - def forward(self, x): - return self.head(F.relu(self.trunk(x))) - - -# ------------------------------------------------------------ -# Data -# ------------------------------------------------------------ -def get_datasets(root): - tf = transforms.Compose([transforms.ToTensor(), - transforms.Normalize((0.1307,), (0.3081,))]) - tr = datasets.MNIST(root, train=True, download=True, transform=tf) - te = datasets.MNIST(root, train=False, download=True, transform=tf) - return tr, te - - -def sample_frames(ds_x, ds_y, B, U, device, rng): - idx = torch.from_numpy(rng.integers(0, ds_x.shape[0], size=(B * U,))) - imgs = ds_x[idx].to(device).view(B, U, 1, 28, 28) - labels = ds_y[idx].to(device).view(B, U) - return imgs, labels - - -def tensorize(ds): - x = torch.stack([img for img, _ in ds]) - y = torch.tensor([lbl for _, lbl in ds]) - return x, y - - -# ------------------------------------------------------------ -# End-to-end forward (mirrors base.semantic_forward) -# ------------------------------------------------------------ -def forward_frames(model, imgs, chan, g_p, g_d, X_pilot, noise_var, fd, delta): - y_emb, m = model.tx(imgs) - X_data = y_emb.to(torch.complex64) - Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(imgs.shape[0], -1), g_p, noise_var) - Y_d = chan.transmit(X_data, g_d, noise_var) - H_ls = Y_p * torch.conj(X_pilot).unsqueeze(0) - rho = base.aging_rho(fd, delta, chan.N, chan.cp) - H_til = (rho / (1.0 + noise_var)) * H_ls - q = 1.0 - (rho ** 2) / (1.0 + noise_var) - Yeq = torch.conj(H_til) * Y_d / (H_til.abs() ** 2 + q + noise_var) - rms = torch.sqrt(torch.mean(Yeq.abs() ** 2, dim=1, keepdim=True) + 1e-12) - return model.rx(Yeq / rms, m) - - -# ------------------------------------------------------------ -# Conventional digital chain (TX classification + QPSK index) -# ------------------------------------------------------------ -def digital_tx(pred_cls, U, N, device): - """4-bit class index -> two QPSK symbols on comb subcarriers.""" - b32 = (pred_cls >> 2) & 3 # (B,U) first 2 bits - b10 = pred_cls & 3 - const = base.qpsk_constellation(device) - s1, s2 = const[b32], const[b10] - B = pred_cls.shape[0] - X = torch.zeros(B, N, dtype=torch.complex64, device=device) - for u in range(U): - ks = torch.arange(u, N, U, device=device) # 16 comb tones - X[:, ks[:8]] = s1[:, u].unsqueeze(1) - X[:, ks[8:]] = s2[:, u].unsqueeze(1) - return X - - -def digital_detect(Y, H_hat, U, N): - const = base.qpsk_constellation(Y.device) - B = Y.shape[0] - out = torch.zeros(B, U, dtype=torch.long, device=Y.device) - for u in range(U): - ks = torch.arange(u, N, U, device=Y.device) - idx_pair = [] - for grp in (ks[:8], ks[8:]): - Z = (torch.conj(H_hat[:, grp]) * Y[:, grp]).sum(dim=1) - metric = torch.real(torch.conj(const).view(1, 4) * Z.unsqueeze(1)) - idx_pair.append(metric.argmax(dim=1)) - out[:, u] = idx_pair[0] * 4 + idx_pair[1] - return out # (B,U) in 0..15 - - -# ------------------------------------------------------------ -# Training / evaluation -# ------------------------------------------------------------ -def train_semantic(args, device, model_cls=None, ckpt="mnist_semantic.pt"): - base.set_seed(args.seed) - tr, _ = get_datasets(args.data_root) - x, y = tensorize(tr) - chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) - X_pilot = base.make_pilot(args.nfft, device) - if model_cls is None: - model_cls = MnistSemanticMA - model = model_cls(args.users, args.dim, args.hidden).to(device) - opt = torch.optim.Adam(model.parameters(), lr=args.lr) - rng = np.random.default_rng(args.seed) - - for step in range(1, args.steps + 1): - snr, fd = base.sample_task(args, rng) - imgs, labels = sample_frames(x, y, args.batch, args.users, device, rng) - g_p, g_d = chan.sample(args.batch, fd, args.delta) - noise_var = 10 ** (-snr / 10.0) - logits = forward_frames(model, imgs, chan, g_p, g_d, X_pilot, - noise_var, fd, args.delta) - loss = F.cross_entropy(logits.reshape(-1, 10), labels.reshape(-1)) - opt.zero_grad(set_to_none=True) - loss.backward() - opt.step() - if step % 200 == 0: - print(f"[mnist-sem {step}/{args.steps}] loss={loss.item():.4f}", flush=True) - - os.makedirs(args.save_dir, exist_ok=True) - torch.save(model.state_dict(), os.path.join(args.save_dir, ckpt)) - print("saved " + ckpt) - - -def train_tx_cls(args, device): - base.set_seed(args.seed) - tr, te = get_datasets(args.data_root) - x, y = tensorize(tr) - xt, yt = tensorize(te) - model = TxClassifier(args.dim).to(device) - opt = torch.optim.Adam(model.parameters(), lr=1e-3) - rng = np.random.default_rng(args.seed) - n = x.shape[0] - for ep in range(3): - perm = torch.from_numpy(rng.permutation(n)) - for i in range(0, n, 512): - idx = perm[i:i + 512] - logits = model(x[idx].to(device)) - loss = F.cross_entropy(logits, y[idx].to(device)) - opt.zero_grad(set_to_none=True) - loss.backward() - opt.step() - with torch.no_grad(): - acc = 0 - for i in range(0, xt.shape[0], 2048): - acc += (model(xt[i:i + 2048].to(device)).argmax(-1) - == yt[i:i + 2048].to(device)).sum().item() - print(f"[tx-cls epoch {ep + 1}] test acc={acc / xt.shape[0]:.4f}", flush=True) - os.makedirs(args.save_dir, exist_ok=True) - torch.save(model.state_dict(), os.path.join(args.save_dir, "mnist_txcls.pt")) - print("saved mnist_txcls.pt") - - - -MNIST_DECODER_KEYS_PREFIX = ("score.", "cls.") - - -def rx_with_fast(model, Yeq, m, fast): - B = Yeq.shape[0] - R = Yeq.unsqueeze(1) * m.unsqueeze(0) - phi = torch.cat([R.real, R.imag], dim=-1) - G = torch.bmm(phi, phi.transpose(1, 2)) / (2 * model.d) - h1 = torch.relu(F.linear(G.reshape(B, -1), fast["score.0.weight"], - fast["score.0.bias"])) - h1 = torch.relu(F.linear(h1, fast["score.2.weight"], - fast["score.2.bias"])) - wsc = F.linear(h1, fast["score.4.weight"], fast["score.4.bias"]) - W = torch.eye(model.U, device=Yeq.device).unsqueeze(0) - W = W + wsc.view(B, model.U, model.U) - z = torch.bmm(W, phi) - h2 = torch.relu(F.linear(z, fast["cls.0.weight"], fast["cls.0.bias"])) - return F.linear(h2, fast["cls.2.weight"], fast["cls.2.bias"]) - - -def forward_frames_fast(model, fast, imgs, chan, g_p, g_d, X_pilot, - noise_var, fd, delta): - y_emb, m = model.tx(imgs) - X_data = y_emb.to(torch.complex64) - Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(imgs.shape[0], -1), - g_p, noise_var) - Y_d = chan.transmit(X_data, g_d, noise_var) - H_ls = Y_p * torch.conj(X_pilot).unsqueeze(0) - rho = base.aging_rho(fd, delta, chan.N, chan.cp) - H_til = (rho / (1.0 + noise_var)) * H_ls - q = 1.0 - (rho ** 2) / (1.0 + noise_var) - Yeq = torch.conj(H_til) * Y_d / (H_til.abs() ** 2 + q + noise_var) - rms = torch.sqrt(torch.mean(Yeq.abs() ** 2, dim=1, keepdim=True) + 1e-12) - return rx_with_fast(model, Yeq / rms, m, fast) - - -def adapt_mnist(model, xs, ys, chan, X_pilot, args, snr, device): - rng = np.random.default_rng(args.seed + 77) - fast = {k: v.detach().clone().requires_grad_(True) - for k, v in model.named_parameters() - if k.startswith(MNIST_DECODER_KEYS_PREFIX)} - for _ in range(args.eval_inner_steps): - imgs, labels = sample_frames(xs, ys, args.support, args.users, - device, rng) - g_p, g_d = chan.sample(args.support, args.eval_fd, args.delta) - noise_var = 10 ** (-snr / 10.0) - with torch.enable_grad(): - logits = forward_frames_fast(model, fast, imgs, chan, g_p, g_d, - X_pilot, noise_var, args.eval_fd, - args.delta) - loss = F.cross_entropy(logits.reshape(-1, 10), - labels.reshape(-1)) - grads = torch.autograd.grad(loss, list(fast.values())) - fast = {k: (p - args.eval_inner_lr * g).detach().requires_grad_(True) - for (k, p), g in zip(fast.items(), grads)} - return {k: v.detach() for k, v in fast.items()} - - -def train_maml_semantic(args, device): - base.set_seed(args.seed) - tr, _ = get_datasets(args.data_root) - x, y = tensorize(tr) - chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) - X_pilot = base.make_pilot(args.nfft, device) - model = MnistSemanticMA(args.users, args.dim, args.hidden).to(device) - warm = os.path.join(args.save_dir, "mnist_semantic.pt") - if os.path.exists(warm): - model.load_state_dict(torch.load(warm, map_location=device)) - print("meta-training warm-started from mnist_semantic.pt", flush=True) - opt = torch.optim.Adam(model.parameters(), lr=args.lr) - rng = np.random.default_rng(args.seed) - - def task_loss(fast, snr, fd): - imgs, labels = sample_frames(x, y, args.batch, args.users, device, - rng) - g_p, g_d = chan.sample(args.batch, fd, args.delta) - noise_var = 10 ** (-snr / 10.0) - logits = forward_frames_fast(model, fast, imgs, chan, g_p, g_d, - X_pilot, noise_var, fd, args.delta) - return F.cross_entropy(logits.reshape(-1, 10), labels.reshape(-1)) - - for step in range(1, args.meta_steps + 1): - meta_loss = 0.0 - for _ in range(args.meta_batch): - snr, fd = base.sample_task(args, rng) - fast = {k: v for k, v in model.named_parameters() - if k.startswith(MNIST_DECODER_KEYS_PREFIX)} - loss_sup = task_loss(fast, snr, fd) - grads = torch.autograd.grad(loss_sup, list(fast.values())) - fast = {k: p - args.inner_lr * g.detach() - for (k, p), g in zip(fast.items(), grads)} - meta_loss = meta_loss + task_loss(fast, snr, fd) - meta_loss = meta_loss / args.meta_batch - opt.zero_grad(set_to_none=True) - meta_loss.backward() - opt.step() - if step % 100 == 0: - print("[mnist-maml %d/%d] loss=%.4f" - % (step, args.meta_steps, meta_loss.item()), flush=True) - os.makedirs(args.save_dir, exist_ok=True) - torch.save(model.state_dict(), - os.path.join(args.save_dir, "mnist_maml.pt")) - print("saved mnist_maml.pt") - - -@torch.no_grad() -def eval_all(args, device): - base.set_seed(args.seed + 3) - _, te = get_datasets(args.data_root) - x, y = tensorize(te) - chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) - X_pilot = base.make_pilot(args.nfft, device) - - sem = MnistSemanticMA(args.users, args.dim, args.hidden).to(device) - sem.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_semantic.pt"), - map_location=device)) - sem.eval() - maml_path = os.path.join(args.save_dir, "mnist_maml.pt") - mm = None - if os.path.exists(maml_path): - mm = MnistSemanticMA(args.users, args.dim, args.hidden).to(device) - mm.load_state_dict(torch.load(maml_path, map_location=device)) - mm.eval() - xs_tr, ys_tr = tensorize(get_datasets(args.data_root)[0]) - tf_path = os.path.join(args.save_dir, "mnist_tf.pt") - tfm = None - if os.path.exists(tf_path): - tfm = MnistTransformerMA(args.users, args.dim, args.hidden).to(device) - tfm.load_state_dict(torch.load(tf_path, map_location=device)) - tfm.eval() - ae_path = os.path.join(args.save_dir, "mnist_ae.pt") - aem = None - if os.path.exists(ae_path): - aem = MnistPerUserAE(args.users, args.dim, args.hidden).to(device) - aem.load_state_dict(torch.load(ae_path, map_location=device)) - aem.eval() - txc = TxClassifier(args.dim).to(device) - txc.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_txcls.pt"), - map_location=device)) - txc.eval() - - rng = np.random.default_rng(args.seed + 3) - csv_path = os.path.join(args.save_dir, "mnist_results.csv") - with open(csv_path, "w", newline="") as f: - w = csv.writer(f) - w.writerow(["snr_db", "fd_norm", "method", "ser"]) - for snr in args.eval_snrs: - noise_var = 10 ** (-snr / 10.0) - errs = {"digital_genie": 0, "digital_pilot": 0, "semantic": 0} - if tfm is not None: - errs["semantic_tf"] = 0 - if aem is not None: - errs["semantic_ae"] = 0 - if mm is not None: - errs["semantic_maml"] = 0 - fast_mm = adapt_mnist(mm, xs_tr, ys_tr, chan, X_pilot, - args, snr, device) - total = 0 - for _ in range(args.eval_nb): - imgs, labels = sample_frames(x, y, args.eval_batch, args.users, - device, rng) - g_p, g_d = chan.sample(args.eval_batch, args.eval_fd, args.delta) - - # proposed semantic chain - logits = forward_frames(sem, imgs, chan, g_p, g_d, X_pilot, - noise_var, args.eval_fd, args.delta) - errs["semantic"] += (logits.argmax(-1) != labels).sum().item() - if tfm is not None: - lg2 = forward_frames(tfm, imgs, chan, g_p, g_d, X_pilot, - noise_var, args.eval_fd, args.delta) - errs["semantic_tf"] += (lg2.argmax(-1) != labels).sum().item() - if aem is not None: - lg4 = forward_frames(aem, imgs, chan, g_p, g_d, X_pilot, - noise_var, args.eval_fd, args.delta) - errs["semantic_ae"] += (lg4.argmax(-1) != labels).sum().item() - if mm is not None: - lg3 = forward_frames_fast(mm, fast_mm, imgs, chan, g_p, - g_d, X_pilot, noise_var, - args.eval_fd, args.delta) - errs["semantic_maml"] += (lg3.argmax(-1) != labels).sum().item() - - # conventional digital chain - B = imgs.shape[0] - pred_cls = txc(imgs.reshape(B * args.users, 1, 28, 28)).argmax(-1) - pred_cls = pred_cls.view(B, args.users) - X_d = digital_tx(pred_cls, args.users, args.nfft, device) - Y_d = chan.transmit(X_d, g_d, noise_var) - Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(B, -1), g_p, noise_var) - H_ls = Y_p * torch.conj(X_pilot).unsqueeze(0) - H_true = chan.genie_H(g_d) - for name, H in [("digital_genie", H_true), ("digital_pilot", H_ls)]: - rec = digital_detect(Y_d, H, args.users, args.nfft) - errs[name] += (rec != labels).sum().item() - total += labels.numel() - - for name, e in errs.items(): - w.writerow([snr, args.eval_fd, name, e / total]) - f.flush() - print(f"snr={snr:5.1f} | " + - " ".join(f"{k}={v / total:.4e}" for k, v in errs.items()), flush=True) - print(f"saved {csv_path}") - - -def make_fig(args): - import matplotlib - matplotlib.use("Agg") - import matplotlib.pyplot as plt - rows = [] - with open(os.path.join(args.save_dir, "mnist_results.csv")) as f: - rows = list(csv.DictReader(f)) - LAB = {"digital_genie": "Digital chain genie CSI", - "digital_pilot": "Digital chain pilot CSI", - "semantic_tf": "Transformer SE separation", - "semantic_ae": "Per-user AE multiple access", - "semantic": "Proposed signed joint", - "semantic_maml": "Proposed signed MAML"} - STY = {"digital_genie": dict(color="gray", marker="^", ls="--"), - "digital_pilot": dict(color="k", marker="v", ls="-"), - "semantic_tf": dict(color="tab:purple", marker="P", ls="-"), - "semantic_ae": dict(color="tab:brown", marker="X", ls="-"), - "semantic": dict(color="tab:red", marker="o", ls="-"), - "semantic_maml": dict(color="tab:green", marker="D", ls="--")} - fig = plt.figure(figsize=(5.2, 3.9)) - ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) - for mkey in LAB: - pts = sorted([(float(r["snr_db"]), float(r["ser"])) - for r in rows if r["method"] == mkey]) - if not pts: - continue - xs, ys = zip(*pts) - ys = [max(v, 1e-5) for v in ys] - ax.semilogy(xs, ys, label=LAB[mkey], ms=4, lw=1.3, **STY[mkey]) - ax.set_xlabel("SNR (dB)") - ax.set_ylabel("SER") - ax.grid(True, which="both", alpha=0.35) - ax.legend(fontsize=7.5, loc="center right", bbox_to_anchor=(0.985, 0.66)) - out = os.path.join(args.fig_dir, f"mnist_ser_vs_snr_fd{args.eval_fd}.pdf") - fig.savefig(out) - print("saved", out) - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--mode", choices=["train", "train-tf", "train-ae", "train-maml", "train-tx-cls", "eval", "fig"], - required=True) - p.add_argument("--users", type=int, default=8) - p.add_argument("--dim", type=int, default=128) - p.add_argument("--hidden", type=int, default=256) - p.add_argument("--nfft", type=int, default=128) - p.add_argument("--cp", type=int, default=16) - p.add_argument("--taps", type=int, default=8) - p.add_argument("--delta", type=int, default=6) - p.add_argument("--train-snrs", type=float, nargs="+", default=[0, 5, 10, 15, 20, 25]) - p.add_argument("--train-fds", type=float, nargs="+", - default=[0.002, 0.005, 0.01, 0.02, 0.05, 0.1]) - p.add_argument("--steps", type=int, default=4000) - p.add_argument("--batch", type=int, default=64) - p.add_argument("--lr", type=float, default=1e-3) - p.add_argument("--meta-steps", type=int, default=2500) - p.add_argument("--meta-batch", type=int, default=4) - p.add_argument("--inner-lr", type=float, default=0.02) - p.add_argument("--support", type=int, default=32) - p.add_argument("--eval-inner-steps", type=int, default=5) - p.add_argument("--eval-inner-lr", type=float, default=0.01) - p.add_argument("--eval-snrs", type=float, nargs="+", - default=[0, 5, 10, 15, 20, 25, 30]) - p.add_argument("--eval-fd", type=float, default=0.05) - p.add_argument("--eval-batch", type=int, default=64) - p.add_argument("--eval-nb", type=int, default=60) - p.add_argument("--seed", type=int, default=0) - p.add_argument("--save-dir", type=str, default="results_mnist") - p.add_argument("--fig-dir", type=str, default="fig") - p.add_argument("--data-root", type=str, default="data_mnist") - args = p.parse_args() - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - print("device:", device) - if args.mode == "train": - train_semantic(args, device) - elif args.mode == "train-maml": - train_maml_semantic(args, device) - elif args.mode == "train-tf": - train_semantic(args, device, model_cls=MnistTransformerMA, ckpt="mnist_tf.pt") - elif args.mode == "train-ae": - train_semantic(args, device, model_cls=MnistPerUserAE, ckpt="mnist_ae.pt") - elif args.mode == "train-tx-cls": - train_tx_cls(args, device) - elif args.mode == "eval": - eval_all(args, device) - elif args.mode == "fig": - make_fig(args) - - -if __name__ == "__main__": - main() +#!/usr/bin/env python3 +# ============================================================ +# c13_mnist.py +# +# Real-data validation: multi-user semantic transmission of MNIST +# images over the time-varying frequency-selective channel of c11. +# +# Proposed: shared CNN semantic encoder -> d-dim embedding per user +# -> learnable user masks + over-the-air superposition -> OFDM +# frame (pilot + data, CSI aging) -> aging-aware LMMSE + RMS norm +# -> signed user-wise attention demux -> shared 10-class head. +# Metric: semantic error rate, P(recovered class != true label). +# +# Conventional scheme: transmit-side classification with the same +# CNN trunk, then digital transmission of the 4-bit class index as +# two QPSK symbols over the user's 16 comb subcarriers (8-fold +# repetition each, MRC with genie or aged pilot CSI). +# +# Usage: +# python3 c13_mnist.py --mode train +# python3 c13_mnist.py --mode train-tx-cls +# python3 c13_mnist.py --mode eval +# python3 c13_mnist.py --mode fig +# ============================================================ + +import argparse +import csv +import math +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision import datasets, transforms + +import c11_doppler_csi as base + + +# ------------------------------------------------------------ +# Models +# ------------------------------------------------------------ +class CNNTrunk(nn.Module): + def __init__(self, out_dim): + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.ReLU(inplace=True), + nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU(inplace=True), + nn.Flatten(), + nn.Linear(32 * 7 * 7, out_dim), + ) + + def forward(self, x): + return self.net(x) + + +class MnistSemanticMA(nn.Module): + """CNN encoder + masks + signed user-wise attention + classifier.""" + + def __init__(self, U, d=128, hidden=256, n_cls=10, score_hidden=64): + super().__init__() + self.U, self.d = U, d + self.encoder = CNNTrunk(d) + self.masks = nn.Parameter(torch.randn(U, d)) + self.score = nn.Sequential( + nn.Linear(U * U, score_hidden), nn.ReLU(inplace=True), + nn.Linear(score_hidden, score_hidden), nn.ReLU(inplace=True), + nn.Linear(score_hidden, U * U), + ) + self.cls = nn.Sequential( + nn.Linear(2 * d, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, n_cls), + ) + + def tx(self, imgs, params=None): + # imgs: (B,U,1,28,28) + B = imgs.shape[0] + e = self.encoder(imgs.reshape(B * self.U, 1, 28, 28)).view(B, self.U, self.d) + m = F.normalize(self.masks, dim=1) + y = (e * m.unsqueeze(0)).sum(dim=1) + y = y / torch.sqrt(torch.mean(y ** 2, dim=1, keepdim=True) + 1e-12) + return y, m + + def rx(self, Yeq, m, params=None): + B = Yeq.shape[0] + R = Yeq.unsqueeze(1) * m.unsqueeze(0) + phi = torch.cat([R.real, R.imag], dim=-1) + T = torch.bmm(phi, phi.transpose(1, 2)) / (2 * self.d) + w = self.score(T.reshape(B, -1)).view(B, self.U, self.U) + W = torch.eye(self.U, device=phi.device).unsqueeze(0) + w + z = torch.bmm(W, phi) + return self.cls(z) # (B,U,10) + + + +class MnistTransformerMA(MnistSemanticMA): + """SOTA variant: Transformer separation on the equalized features.""" + + def __init__(self, U, d=128, hidden=256, n_cls=10, n_heads=4, n_layers=2): + super().__init__(U, d, hidden, n_cls) + import torch.nn as nn + layer = nn.TransformerEncoderLayer(d_model=hidden, nhead=n_heads, + dim_feedforward=2 * hidden, + batch_first=True) + self.inp = nn.Linear(2 * d, hidden) + self.sep = nn.TransformerEncoder(layer, num_layers=n_layers) + self.out = nn.Linear(hidden, n_cls) + + def rx(self, Yeq, m, params=None): + import torch + R = Yeq.unsqueeze(1) * m.unsqueeze(0) + phi = torch.cat([R.real, R.imag], dim=-1) + return self.out(self.sep(self.inp(phi))) + + +class MnistPerUserAE(nn.Module): + """DeepMA-style per-user AE: dedicated per-user CNN encoders (no + shared masks) and dedicated per-user decoder heads, superposed on + the identical full-band resource with unit transmit power.""" + + def __init__(self, U, d=128, hidden=256, n_cls=10): + super().__init__() + self.U, self.d = U, d + self.encs = nn.ModuleList([CNNTrunk(d) for _ in range(U)]) + self.heads = nn.ModuleList([ + nn.Sequential(nn.Linear(2 * d, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, n_cls)) + for _ in range(U) + ]) + + def tx(self, imgs, params=None): + e = torch.stack([self.encs[u](imgs[:, u]) for u in range(self.U)], + dim=1) # (B,U,d) + y = e.sum(dim=1) + y = y / torch.sqrt(torch.mean(y ** 2, dim=1, keepdim=True) + 1e-12) + return y, None + + def rx(self, Yeq, m, params=None): + phi = torch.cat([Yeq.real, Yeq.imag], dim=-1) # (B,2d) + return torch.stack([h(phi) for h in self.heads], dim=1) + + +class TxClassifier(nn.Module): + """Transmit-side classifier for the conventional digital chain.""" + + def __init__(self, d=128, n_cls=10): + super().__init__() + self.trunk = CNNTrunk(d) + self.head = nn.Linear(d, n_cls) + + def forward(self, x): + return self.head(F.relu(self.trunk(x))) + + +# ------------------------------------------------------------ +# Data +# ------------------------------------------------------------ +def get_datasets(root): + tf = transforms.Compose([transforms.ToTensor(), + transforms.Normalize((0.1307,), (0.3081,))]) + tr = datasets.MNIST(root, train=True, download=True, transform=tf) + te = datasets.MNIST(root, train=False, download=True, transform=tf) + return tr, te + + +def sample_frames(ds_x, ds_y, B, U, device, rng): + idx = torch.from_numpy(rng.integers(0, ds_x.shape[0], size=(B * U,))) + imgs = ds_x[idx].to(device).view(B, U, 1, 28, 28) + labels = ds_y[idx].to(device).view(B, U) + return imgs, labels + + +def tensorize(ds): + x = torch.stack([img for img, _ in ds]) + y = torch.tensor([lbl for _, lbl in ds]) + return x, y + + +# ------------------------------------------------------------ +# End-to-end forward (mirrors base.semantic_forward) +# ------------------------------------------------------------ +def forward_frames(model, imgs, chan, g_p, g_d, X_pilot, noise_var, fd, delta): + y_emb, m = model.tx(imgs) + X_data = y_emb.to(torch.complex64) + Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(imgs.shape[0], -1), g_p, noise_var) + Y_d = chan.transmit(X_data, g_d, noise_var) + H_ls = Y_p * torch.conj(X_pilot).unsqueeze(0) + rho = base.aging_rho(fd, delta, chan.N, chan.cp) + H_til = (rho / (1.0 + noise_var)) * H_ls + q = 1.0 - (rho ** 2) / (1.0 + noise_var) + Yeq = torch.conj(H_til) * Y_d / (H_til.abs() ** 2 + q + noise_var) + rms = torch.sqrt(torch.mean(Yeq.abs() ** 2, dim=1, keepdim=True) + 1e-12) + return model.rx(Yeq / rms, m) + + +# ------------------------------------------------------------ +# Conventional digital chain (TX classification + QPSK index) +# ------------------------------------------------------------ +def digital_tx(pred_cls, U, N, device): + """4-bit class index -> two QPSK symbols on comb subcarriers.""" + b32 = (pred_cls >> 2) & 3 # (B,U) first 2 bits + b10 = pred_cls & 3 + const = base.qpsk_constellation(device) + s1, s2 = const[b32], const[b10] + B = pred_cls.shape[0] + X = torch.zeros(B, N, dtype=torch.complex64, device=device) + for u in range(U): + ks = torch.arange(u, N, U, device=device) # 16 comb tones + X[:, ks[:8]] = s1[:, u].unsqueeze(1) + X[:, ks[8:]] = s2[:, u].unsqueeze(1) + return X + + +def digital_detect(Y, H_hat, U, N): + const = base.qpsk_constellation(Y.device) + B = Y.shape[0] + out = torch.zeros(B, U, dtype=torch.long, device=Y.device) + for u in range(U): + ks = torch.arange(u, N, U, device=Y.device) + idx_pair = [] + for grp in (ks[:8], ks[8:]): + Z = (torch.conj(H_hat[:, grp]) * Y[:, grp]).sum(dim=1) + metric = torch.real(torch.conj(const).view(1, 4) * Z.unsqueeze(1)) + idx_pair.append(metric.argmax(dim=1)) + out[:, u] = idx_pair[0] * 4 + idx_pair[1] + return out # (B,U) in 0..15 + + +# ------------------------------------------------------------ +# Training / evaluation +# ------------------------------------------------------------ +def train_semantic(args, device, model_cls=None, ckpt="mnist_semantic.pt"): + base.set_seed(args.seed) + tr, _ = get_datasets(args.data_root) + x, y = tensorize(tr) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + if model_cls is None: + model_cls = MnistSemanticMA + model = model_cls(args.users, args.dim, args.hidden).to(device) + opt = torch.optim.Adam(model.parameters(), lr=args.lr) + rng = np.random.default_rng(args.seed) + + for step in range(1, args.steps + 1): + snr, fd = base.sample_task(args, rng) + imgs, labels = sample_frames(x, y, args.batch, args.users, device, rng) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = forward_frames(model, imgs, chan, g_p, g_d, X_pilot, + noise_var, fd, args.delta) + loss = F.cross_entropy(logits.reshape(-1, 10), labels.reshape(-1)) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + if step % 200 == 0: + print(f"[mnist-sem {step}/{args.steps}] loss={loss.item():.4f}", flush=True) + + os.makedirs(args.save_dir, exist_ok=True) + torch.save(model.state_dict(), os.path.join(args.save_dir, ckpt)) + print("saved " + ckpt) + + +def train_tx_cls(args, device): + base.set_seed(args.seed) + tr, te = get_datasets(args.data_root) + x, y = tensorize(tr) + xt, yt = tensorize(te) + model = TxClassifier(args.dim).to(device) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + rng = np.random.default_rng(args.seed) + n = x.shape[0] + for ep in range(3): + perm = torch.from_numpy(rng.permutation(n)) + for i in range(0, n, 512): + idx = perm[i:i + 512] + logits = model(x[idx].to(device)) + loss = F.cross_entropy(logits, y[idx].to(device)) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + with torch.no_grad(): + acc = 0 + for i in range(0, xt.shape[0], 2048): + acc += (model(xt[i:i + 2048].to(device)).argmax(-1) + == yt[i:i + 2048].to(device)).sum().item() + print(f"[tx-cls epoch {ep + 1}] test acc={acc / xt.shape[0]:.4f}", flush=True) + os.makedirs(args.save_dir, exist_ok=True) + torch.save(model.state_dict(), os.path.join(args.save_dir, "mnist_txcls.pt")) + print("saved mnist_txcls.pt") + + + +MNIST_DECODER_KEYS_PREFIX = ("score.", "cls.") + + +def rx_with_fast(model, Yeq, m, fast): + B = Yeq.shape[0] + R = Yeq.unsqueeze(1) * m.unsqueeze(0) + phi = torch.cat([R.real, R.imag], dim=-1) + G = torch.bmm(phi, phi.transpose(1, 2)) / (2 * model.d) + h1 = torch.relu(F.linear(G.reshape(B, -1), fast["score.0.weight"], + fast["score.0.bias"])) + h1 = torch.relu(F.linear(h1, fast["score.2.weight"], + fast["score.2.bias"])) + wsc = F.linear(h1, fast["score.4.weight"], fast["score.4.bias"]) + W = torch.eye(model.U, device=Yeq.device).unsqueeze(0) + W = W + wsc.view(B, model.U, model.U) + z = torch.bmm(W, phi) + h2 = torch.relu(F.linear(z, fast["cls.0.weight"], fast["cls.0.bias"])) + return F.linear(h2, fast["cls.2.weight"], fast["cls.2.bias"]) + + +def forward_frames_fast(model, fast, imgs, chan, g_p, g_d, X_pilot, + noise_var, fd, delta): + y_emb, m = model.tx(imgs) + X_data = y_emb.to(torch.complex64) + Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(imgs.shape[0], -1), + g_p, noise_var) + Y_d = chan.transmit(X_data, g_d, noise_var) + H_ls = Y_p * torch.conj(X_pilot).unsqueeze(0) + rho = base.aging_rho(fd, delta, chan.N, chan.cp) + H_til = (rho / (1.0 + noise_var)) * H_ls + q = 1.0 - (rho ** 2) / (1.0 + noise_var) + Yeq = torch.conj(H_til) * Y_d / (H_til.abs() ** 2 + q + noise_var) + rms = torch.sqrt(torch.mean(Yeq.abs() ** 2, dim=1, keepdim=True) + 1e-12) + return rx_with_fast(model, Yeq / rms, m, fast) + + +def adapt_mnist(model, xs, ys, chan, X_pilot, args, snr, device): + rng = np.random.default_rng(args.seed + 77) + fast = {k: v.detach().clone().requires_grad_(True) + for k, v in model.named_parameters() + if k.startswith(MNIST_DECODER_KEYS_PREFIX)} + for _ in range(args.eval_inner_steps): + imgs, labels = sample_frames(xs, ys, args.support, args.users, + device, rng) + g_p, g_d = chan.sample(args.support, args.eval_fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + with torch.enable_grad(): + logits = forward_frames_fast(model, fast, imgs, chan, g_p, g_d, + X_pilot, noise_var, args.eval_fd, + args.delta) + loss = F.cross_entropy(logits.reshape(-1, 10), + labels.reshape(-1)) + grads = torch.autograd.grad(loss, list(fast.values())) + fast = {k: (p - args.eval_inner_lr * g).detach().requires_grad_(True) + for (k, p), g in zip(fast.items(), grads)} + return {k: v.detach() for k, v in fast.items()} + + +def train_maml_semantic(args, device): + base.set_seed(args.seed) + tr, _ = get_datasets(args.data_root) + x, y = tensorize(tr) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + model = MnistSemanticMA(args.users, args.dim, args.hidden).to(device) + warm = os.path.join(args.save_dir, "mnist_semantic.pt") + if os.path.exists(warm): + model.load_state_dict(torch.load(warm, map_location=device)) + print("meta-training warm-started from mnist_semantic.pt", flush=True) + opt = torch.optim.Adam(model.parameters(), lr=args.lr) + rng = np.random.default_rng(args.seed) + + def task_loss(fast, snr, fd): + imgs, labels = sample_frames(x, y, args.batch, args.users, device, + rng) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = forward_frames_fast(model, fast, imgs, chan, g_p, g_d, + X_pilot, noise_var, fd, args.delta) + return F.cross_entropy(logits.reshape(-1, 10), labels.reshape(-1)) + + for step in range(1, args.meta_steps + 1): + meta_loss = 0.0 + for _ in range(args.meta_batch): + snr, fd = base.sample_task(args, rng) + fast = {k: v for k, v in model.named_parameters() + if k.startswith(MNIST_DECODER_KEYS_PREFIX)} + loss_sup = task_loss(fast, snr, fd) + grads = torch.autograd.grad(loss_sup, list(fast.values())) + fast = {k: p - args.inner_lr * g.detach() + for (k, p), g in zip(fast.items(), grads)} + meta_loss = meta_loss + task_loss(fast, snr, fd) + meta_loss = meta_loss / args.meta_batch + opt.zero_grad(set_to_none=True) + meta_loss.backward() + opt.step() + if step % 100 == 0: + print("[mnist-maml %d/%d] loss=%.4f" + % (step, args.meta_steps, meta_loss.item()), flush=True) + os.makedirs(args.save_dir, exist_ok=True) + torch.save(model.state_dict(), + os.path.join(args.save_dir, "mnist_maml.pt")) + print("saved mnist_maml.pt") + + +@torch.no_grad() +def eval_all(args, device): + base.set_seed(args.seed + 3) + _, te = get_datasets(args.data_root) + x, y = tensorize(te) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + + sem = MnistSemanticMA(args.users, args.dim, args.hidden).to(device) + sem.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_semantic.pt"), + map_location=device)) + sem.eval() + maml_path = os.path.join(args.save_dir, "mnist_maml.pt") + mm = None + if os.path.exists(maml_path): + mm = MnistSemanticMA(args.users, args.dim, args.hidden).to(device) + mm.load_state_dict(torch.load(maml_path, map_location=device)) + mm.eval() + xs_tr, ys_tr = tensorize(get_datasets(args.data_root)[0]) + tf_path = os.path.join(args.save_dir, "mnist_tf.pt") + tfm = None + if os.path.exists(tf_path): + tfm = MnistTransformerMA(args.users, args.dim, args.hidden).to(device) + tfm.load_state_dict(torch.load(tf_path, map_location=device)) + tfm.eval() + ae_path = os.path.join(args.save_dir, "mnist_ae.pt") + aem = None + if os.path.exists(ae_path): + aem = MnistPerUserAE(args.users, args.dim, args.hidden).to(device) + aem.load_state_dict(torch.load(ae_path, map_location=device)) + aem.eval() + txc = TxClassifier(args.dim).to(device) + txc.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_txcls.pt"), + map_location=device)) + txc.eval() + + rng = np.random.default_rng(args.seed + 3) + csv_path = os.path.join(args.save_dir, "mnist_results.csv") + with open(csv_path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["snr_db", "fd_norm", "method", "ser"]) + for snr in args.eval_snrs: + noise_var = 10 ** (-snr / 10.0) + errs = {"digital_genie": 0, "digital_pilot": 0, "semantic": 0} + if tfm is not None: + errs["semantic_tf"] = 0 + if aem is not None: + errs["semantic_ae"] = 0 + if mm is not None: + errs["semantic_maml"] = 0 + fast_mm = adapt_mnist(mm, xs_tr, ys_tr, chan, X_pilot, + args, snr, device) + total = 0 + for _ in range(args.eval_nb): + imgs, labels = sample_frames(x, y, args.eval_batch, args.users, + device, rng) + g_p, g_d = chan.sample(args.eval_batch, args.eval_fd, args.delta) + + # proposed semantic chain + logits = forward_frames(sem, imgs, chan, g_p, g_d, X_pilot, + noise_var, args.eval_fd, args.delta) + errs["semantic"] += (logits.argmax(-1) != labels).sum().item() + if tfm is not None: + lg2 = forward_frames(tfm, imgs, chan, g_p, g_d, X_pilot, + noise_var, args.eval_fd, args.delta) + errs["semantic_tf"] += (lg2.argmax(-1) != labels).sum().item() + if aem is not None: + lg4 = forward_frames(aem, imgs, chan, g_p, g_d, X_pilot, + noise_var, args.eval_fd, args.delta) + errs["semantic_ae"] += (lg4.argmax(-1) != labels).sum().item() + if mm is not None: + lg3 = forward_frames_fast(mm, fast_mm, imgs, chan, g_p, + g_d, X_pilot, noise_var, + args.eval_fd, args.delta) + errs["semantic_maml"] += (lg3.argmax(-1) != labels).sum().item() + + # conventional digital chain + B = imgs.shape[0] + pred_cls = txc(imgs.reshape(B * args.users, 1, 28, 28)).argmax(-1) + pred_cls = pred_cls.view(B, args.users) + X_d = digital_tx(pred_cls, args.users, args.nfft, device) + Y_d = chan.transmit(X_d, g_d, noise_var) + Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(B, -1), g_p, noise_var) + H_ls = Y_p * torch.conj(X_pilot).unsqueeze(0) + H_true = chan.genie_H(g_d) + for name, H in [("digital_genie", H_true), ("digital_pilot", H_ls)]: + rec = digital_detect(Y_d, H, args.users, args.nfft) + errs[name] += (rec != labels).sum().item() + total += labels.numel() + + for name, e in errs.items(): + w.writerow([snr, args.eval_fd, name, e / total]) + f.flush() + print(f"snr={snr:5.1f} | " + + " ".join(f"{k}={v / total:.4e}" for k, v in errs.items()), flush=True) + print(f"saved {csv_path}") + + +def make_fig(args): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + rows = [] + with open(os.path.join(args.save_dir, "mnist_results.csv")) as f: + rows = list(csv.DictReader(f)) + LAB = {"digital_genie": "Digital chain genie CSI", + "digital_pilot": "Digital chain pilot CSI", + "semantic_tf": "Transformer SE separation", + "semantic_ae": "Per-user AE", + "semantic": "Proposed signed joint", + "semantic_maml": "Proposed signed MAML"} + STY = {"digital_genie": dict(color="gray", marker="^", ls="--"), + "digital_pilot": dict(color="k", marker="v", ls="-"), + "semantic_tf": dict(color="tab:purple", marker="P", ls="-"), + "semantic_ae": dict(color="tab:brown", marker="X", ls="-"), + "semantic": dict(color="tab:red", marker="o", ls="-"), + "semantic_maml": dict(color="tab:green", marker="D", ls="--")} + plt.rcParams.update({"font.size": 13, "axes.labelsize": 13, + "xtick.labelsize": 12, "ytick.labelsize": 12, + "axes.linewidth": 1.1, "grid.linewidth": 0.8, + "xtick.major.width": 1.1, "ytick.major.width": 1.1, + "xtick.minor.width": 0.8, "ytick.minor.width": 0.8, + "xtick.major.size": 4.5, "ytick.major.size": 4.5}) + fig = plt.figure(figsize=(5.2, 3.9)) + ax = fig.add_axes([0.155, 0.145, 0.82, 0.82]) + for mkey in LAB: + pts = sorted([(float(r["snr_db"]), float(r["ser"])) + for r in rows if r["method"] == mkey]) + if not pts: + continue + xs, ys = zip(*pts) + ys = [max(v, 1e-5) for v in ys] + ax.semilogy(xs, ys, label=LAB[mkey], ms=5, lw=1.8, **STY[mkey]) + ax.set_xlabel("SNR (dB)") + ax.set_ylabel("SER") + ax.grid(True, which="both", alpha=0.35) + ax.legend(fontsize=9, framealpha=1.0, labelspacing=0.3, + handlelength=1.8, loc="center right", bbox_to_anchor=(0.985, 0.66)) + out = os.path.join(args.fig_dir, f"mnist_ser_vs_snr_fd{args.eval_fd}.pdf") + fig.savefig(out) + print("saved", out) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["train", "train-tf", "train-ae", "train-maml", "train-tx-cls", "eval", "fig"], + required=True) + p.add_argument("--users", type=int, default=8) + p.add_argument("--dim", type=int, default=128) + p.add_argument("--hidden", type=int, default=256) + p.add_argument("--nfft", type=int, default=128) + p.add_argument("--cp", type=int, default=16) + p.add_argument("--taps", type=int, default=8) + p.add_argument("--delta", type=int, default=6) + p.add_argument("--train-snrs", type=float, nargs="+", default=[0, 5, 10, 15, 20, 25]) + p.add_argument("--train-fds", type=float, nargs="+", + default=[0.002, 0.005, 0.01, 0.02, 0.05, 0.1]) + p.add_argument("--steps", type=int, default=4000) + p.add_argument("--batch", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--meta-steps", type=int, default=2500) + p.add_argument("--meta-batch", type=int, default=4) + p.add_argument("--inner-lr", type=float, default=0.02) + p.add_argument("--support", type=int, default=32) + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.01) + p.add_argument("--eval-snrs", type=float, nargs="+", + default=[0, 5, 10, 15, 20, 25, 30]) + p.add_argument("--eval-fd", type=float, default=0.05) + p.add_argument("--eval-batch", type=int, default=64) + p.add_argument("--eval-nb", type=int, default=60) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--save-dir", type=str, default="results_mnist") + p.add_argument("--fig-dir", type=str, default="fig") + p.add_argument("--data-root", type=str, default="data_mnist") + args = p.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("device:", device) + if args.mode == "train": + train_semantic(args, device) + elif args.mode == "train-maml": + train_maml_semantic(args, device) + elif args.mode == "train-tf": + train_semantic(args, device, model_cls=MnistTransformerMA, ckpt="mnist_tf.pt") + elif args.mode == "train-ae": + train_semantic(args, device, model_cls=MnistPerUserAE, ckpt="mnist_ae.pt") + elif args.mode == "train-tx-cls": + train_tx_cls(args, device) + elif args.mode == "eval": + eval_all(args, device) + elif args.mode == "fig": + make_fig(args) + + +if __name__ == "__main__": + main() diff --git a/c18_mnist_doppler.py b/c18_mnist_doppler.py index 3108d6e..d080ad2 100755 --- a/c18_mnist_doppler.py +++ b/c18_mnist_doppler.py @@ -114,7 +114,7 @@ def make_fig(args): LAB = {"digital_genie": "Digital chain genie CSI", "digital_pilot": "Digital chain pilot CSI", "semantic_tf": "Transformer SE separation", - "semantic_ae": "Per-user AE multiple access", + "semantic_ae": "Per-user AE", "semantic": "Proposed signed joint", "semantic_maml": "Proposed signed MAML"} STY = {"digital_genie": dict(color="gray", marker="^", ls="--"), @@ -123,8 +123,14 @@ def make_fig(args): "semantic_ae": dict(color="tab:brown", marker="X", ls="-"), "semantic": dict(color="tab:red", marker="o", ls="-"), "semantic_maml": dict(color="tab:green", marker="D", ls="--")} + plt.rcParams.update({"font.size": 13, "axes.labelsize": 13, + "xtick.labelsize": 12, "ytick.labelsize": 12, + "axes.linewidth": 1.1, "grid.linewidth": 0.8, + "xtick.major.width": 1.1, "ytick.major.width": 1.1, + "xtick.minor.width": 0.8, "ytick.minor.width": 0.8, + "xtick.major.size": 4.5, "ytick.major.size": 4.5}) fig = plt.figure(figsize=(5.2, 3.9)) - ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) + ax = fig.add_axes([0.155, 0.145, 0.82, 0.82]) for mkey in LAB: pts = sorted([(float(r["fd_norm"]), float(r["ser"])) for r in rows if r["method"] == mkey]) @@ -132,11 +138,12 @@ def make_fig(args): if not pts: continue xs, ys = zip(*pts) - ax.semilogy(xs, ys, label=LAB[mkey], ms=4, lw=1.3, **STY[mkey]) + ax.semilogy(xs, ys, label=LAB[mkey], ms=5, lw=1.8, **STY[mkey]) ax.set_xlabel(r"Normalized Doppler $f_D T_{\mathrm{sym}}$") ax.set_ylabel("SER") ax.grid(True, which="both", alpha=0.35) - ax.legend(fontsize=7.5, loc="center right", bbox_to_anchor=(0.985, 0.72)) + ax.legend(fontsize=9, framealpha=1.0, labelspacing=0.3, + handlelength=1.8, loc="center right", bbox_to_anchor=(0.985, 0.72)) out = os.path.join(args.fig_dir, f"mnist_ser_vs_doppler_snr{int(args.eval_snr)}.pdf") fig.savefig(out) diff --git a/c19_mnist_epoch.py b/c19_mnist_epoch.py index 2586c23..81b3bd2 100755 --- a/c19_mnist_epoch.py +++ b/c19_mnist_epoch.py @@ -22,7 +22,7 @@ import c13_mnist as m13 MODELS = { "semantic": (m13.MnistSemanticMA, "Proposed signed joint"), "semantic_tf": (m13.MnistTransformerMA, "Transformer SE separation"), - "semantic_ae": (m13.MnistPerUserAE, "Per-user AE multiple access"), + "semantic_ae": (m13.MnistPerUserAE, "Per-user AE"), } @@ -202,24 +202,30 @@ def main(): STY = {"semantic": dict(color="tab:red", marker="o", ls="-"), "semantic_tf": dict(color="tab:purple", marker="P", ls="-"), "semantic_ae": dict(color="tab:brown", marker="X", ls="-")} + plt.rcParams.update({"font.size": 13, "axes.labelsize": 13, + "xtick.labelsize": 12, "ytick.labelsize": 12, + "axes.linewidth": 1.1, "grid.linewidth": 0.8, + "xtick.major.width": 1.1, "ytick.major.width": 1.1, + "xtick.minor.width": 0.8, "ytick.minor.width": 0.8, + "xtick.major.size": 4.5, "ytick.major.size": 4.5}) fig = plt.figure(figsize=(5.2, 3.9)) - ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) + ax = fig.add_axes([0.155, 0.145, 0.82, 0.82]) xgrid = sorted({int(r["step"]) for r in rows if r["method"] == "semantic"}) if "digital_genie" in floors: ax.semilogy(xgrid, [floors["digital_genie"]] * len(xgrid), - label="Digital chain genie CSI", ms=3.5, lw=1.2, + label="Digital chain genie CSI", ms=5, lw=1.7, color="gray", marker="^", ls="--") if "digital_pilot" in floors: ax.semilogy(xgrid, [floors["digital_pilot"]] * len(xgrid), - label="Digital chain pilot CSI", ms=3.5, lw=1.2, + label="Digital chain pilot CSI", ms=5, lw=1.7, color="k", marker="v", ls="-") for key in ["semantic_tf", "semantic_ae", "semantic"]: label = MODELS[key][1] pts = sorted([(int(r["step"]), float(r["ser"])) for r in rows if r["method"] == key]) xs, ys = zip(*pts) - ax.semilogy(xs, ys, label=label, ms=3.5, lw=1.3, **STY[key]) + ax.semilogy(xs, ys, label=label, ms=5, lw=1.8, **STY[key]) # signed MAML: the deployed receiver applies the five-step # task-conditional adaptation at every checkpoint. During the # warm-start phase the adapted SER of the evolving joint model is @@ -239,17 +245,29 @@ def main(): pre_pts = sorted([(int(r["step"]), float(r["ser"])) for r in rows if r["method"] == "semantic" and int(r["step"]) < warm_end]) - pts = pre_pts + meta_pts - xs, ys = zip(*pts) - ax.semilogy(xs, ys, label="Proposed signed MAML", ms=3.5, lw=1.3, - color="tab:green", marker="D", ls="--") - ax.axvline(warm_end, color="gray", ls=":", lw=1.0) + # Left of the dotted line no meta-training has happened yet, so + # the curve is the joint model with the same five-step + # adaptation applied. It is drawn dashed with open markers and + # carries its own legend entry, since it is the control + # condition rather than the proposed MAML receiver. + pre_seg = pre_pts + meta_pts[:1] + xs, ys = zip(*pre_seg) + ax.semilogy(xs, ys, label="Adaptation from joint model", ms=5, + lw=1.8, color="tab:green", marker="D", ls="--", + markerfacecolor="none") + xs, ys = zip(*meta_pts) + ax.semilogy(xs, ys, label="Proposed signed MAML", ms=5, lw=1.8, + color="tab:green", marker="D", ls="-") + ax.axvline(warm_end, color="gray", ls=":", lw=1.4) ax.set_xlabel("Training step") ax.set_ylabel("SER") if "digital_genie" in floors: + # extra headroom below the genie floor so that the seven-entry + # legend sits in free space instead of over the curves ax.set_ylim(bottom=floors["digital_genie"] * 0.5) ax.grid(True, which="both", alpha=0.35) - ax.legend(fontsize=7.5, loc="center left", bbox_to_anchor=(0.02, 0.32)) + ax.legend(fontsize=9, loc="center left", bbox_to_anchor=(0.02, 0.34), + framealpha=1.0, labelspacing=0.3, handlelength=1.8) out = os.path.join(args.fig_dir, "mnist_epoch_convergence.pdf") fig.savefig(out) print("saved", out) diff --git a/c20_bert.py b/c20_bert.py index fe18a37..c7b33a2 100755 --- a/c20_bert.py +++ b/c20_bert.py @@ -441,7 +441,7 @@ def make_fig(args): LAB = {"digital_genie": "Digital chain genie CSI", "digital_pilot": "Digital chain pilot CSI", "semantic_tf": "Transformer SE separation", - "semantic_ae": "Per-user AE multiple access", + "semantic_ae": "Per-user AE", "semantic": "Proposed signed joint", "semantic_maml": "Proposed signed MAML"} STY = {"digital_genie": dict(color="gray", marker="^", ls="--"), @@ -450,8 +450,14 @@ def make_fig(args): "semantic_ae": dict(color="tab:brown", marker="X", ls="-"), "semantic": dict(color="tab:red", marker="o", ls="-"), "semantic_maml": dict(color="tab:green", marker="D", ls="--")} + plt.rcParams.update({"font.size": 13, "axes.labelsize": 13, + "xtick.labelsize": 12, "ytick.labelsize": 12, + "axes.linewidth": 1.1, "grid.linewidth": 0.8, + "xtick.major.width": 1.1, "ytick.major.width": 1.1, + "xtick.minor.width": 0.8, "ytick.minor.width": 0.8, + "xtick.major.size": 4.5, "ytick.major.size": 4.5}) fig = plt.figure(figsize=(5.2, 3.9)) - ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) + ax = fig.add_axes([0.155, 0.145, 0.82, 0.82]) for mkey in LAB: pts = sorted([(float(r["snr_db"]), float(r["ser"])) for r in rows if r["method"] == mkey]) @@ -459,11 +465,12 @@ def make_fig(args): if not pts: continue xs, ys = zip(*pts) - ax.semilogy(xs, ys, label=LAB[mkey], ms=4, lw=1.3, **STY[mkey]) + ax.semilogy(xs, ys, label=LAB[mkey], ms=5, lw=1.8, **STY[mkey]) ax.set_xlabel("SNR (dB)") ax.set_ylabel("SER") ax.grid(True, which="both", alpha=0.35) - ax.legend(fontsize=7.5, loc="center right", bbox_to_anchor=(0.985, 0.66)) + ax.legend(fontsize=9, framealpha=1.0, labelspacing=0.3, + handlelength=1.8, loc="center right", bbox_to_anchor=(0.985, 0.66)) out = os.path.join(args.fig_dir, f"bert_ser_vs_snr_fd{args.eval_fd}.pdf") fig.savefig(out) diff --git a/fig_replot_complexity.py b/fig_replot_complexity.py new file mode 100755 index 0000000..a2840a2 --- /dev/null +++ b/fig_replot_complexity.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +# Replot Fig. 4: demultiplexing-stage complexity ratio vs U/T, +# now including the signed user-wise attention variant. +# +# O_transformer = L * T * d * dff (FFN-dominated, L = 12 layers) +# O_softmax = U * d^2 + U^2 * d (key/value projections + scores) +# O_signed = U^2 * d + U^2 * 2h + h^2 (Gram + score network, h = 64) +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +T, L, d, h = 32, 12, 128, 64 +U = np.arange(16, 65) +x = U / T + +plt.rcParams.update({"font.size": 13, "axes.labelsize": 13, + "xtick.labelsize": 12, "ytick.labelsize": 12, + "axes.linewidth": 1.1, "grid.linewidth": 0.8, + "xtick.major.width": 1.1, "ytick.major.width": 1.1, + "xtick.minor.width": 0.8, "ytick.minor.width": 0.8, + "xtick.major.size": 4.5, "ytick.major.size": 4.5}) +fig = plt.figure(figsize=(5.2, 3.9)) +ax = fig.add_axes([0.155, 0.145, 0.82, 0.82]) +colors = {1.0: "tab:blue", 0.5: "tab:orange", 0.25: "tab:green"} +for r in [1.0, 0.5, 0.25]: # r = d / dff + dff = d / r + o_tf = L * T * d * dff + o_soft = U * d ** 2 + U ** 2 * d + o_sgn = U ** 2 * d + U ** 2 * 2 * h + h ** 2 + ax.plot(x, o_soft / o_tf, color=colors[r], ls="-", lw=1.9, + label=f"Softmax, $d/d_{{\\mathrm{{ff}}}}$={r:g}") + ax.plot(x, o_sgn / o_tf, color=colors[r], ls="--", lw=1.9, + label=f"Signed, $d/d_{{\\mathrm{{ff}}}}$={r:g}") +ax.set_yscale("log") +ax.set_xlabel(r"User-to-token ratio $U/T$") +ax.set_ylabel(r"$\mathcal{O}_{\mathrm{Attention}}/\mathcal{O}_{\mathrm{Transformer}}$") +ax.grid(True, which="both", alpha=0.35) +ax.legend(fontsize=9, ncol=2, loc="lower right") +fig.savefig("fig/complexity_ratio_vs_UT_dff.pdf") +print("saved fig/complexity_ratio_vs_UT_dff.pdf")