commit bed475c954a60916fa5455a322e020bb534da141 Author: KiHoLee Date: Sun Aug 2 16:49:49 2026 +0900 Code and stored results for AI-native multi-user semantic communications (JSAC submission) diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..eee60b2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +data_mnist/ +results_bert/bert_feats.pt +__pycache__/ +*.pyc diff --git a/README.md b/README.md new file mode 100755 index 0000000..bbc74e9 --- /dev/null +++ b/README.md @@ -0,0 +1,50 @@ +# JSAC_AIRAN + +Code and stored results for the manuscript +**"AI-Native Multi-User Semantic Communications via Meta-Learned Attention for 6G AI-RAN"** +(submitted to IEEE Journal on Selected Areas in Communications, +Special Issue on Towards Open and Intelligent 6G RAN). + +All experiments use fixed random seeds, and every figure in the paper is +regenerated by a single script from the stored CSV results in this +repository, without rerunning the experiments. + +## Environment + +- WSL2 Ubuntu, Python 3.14, PyTorch with CUDA (experiments ran on an + NVIDIA GPU; CPU fallback works for plotting and evaluation) +- `transformers` and `datasets` (BERT feature extraction only) +- MNIST downloads automatically via `torchvision`; AG News via + `datasets` (`fancyzhx/ag_news`) + +## Figure and table map + +| Paper item | Regenerate figure (from stored CSV) | Full rerun | +|---|---|---| +| Table III (MNIST, flat Rayleigh) | values in `results_mnist/mnist_flat.csv` | `python3 c21_mnist_flat.py --mode run` | +| Fig. 5 (convergence + MAML meta phase) | `python3 c19_mnist_epoch.py --mode fig` | `c19 --mode run`, `c19 --mode run-maml --steps 7500`, `c22_maml_epoch.py --pretrain-steps 7500` | +| Fig. 6 (MNIST SER vs SNR) | `python3 c13_mnist.py --mode fig` | `c13 --mode train / train-tf / train-ae / train-tx-cls / train-maml / eval` | +| Fig. 7 (MNIST SER vs Doppler) | `python3 c18_mnist_doppler.py --mode fig` | `c18 --mode run` | +| Fig. 8 (BERT text SER vs SNR) | `python3 c20_bert.py --mode fig` | `c20 --mode cache / train-tx-cls / train / train-tf / train-ae / train-maml / eval` | +| Table IV (summary) | values in `results_mnist/mnist_results.csv`, `mnist_doppler.csv` | see Fig. 6 and Fig. 7 rows | + +Run every command from the repository root. Trained checkpoints +(`*.pt`) are included, so `--mode eval` and `--mode fig` work without +retraining. The BERT feature cache (`results_bert/bert_feats.pt`, +about 85 MB) is excluded and is regenerated by +`python3 c20_bert.py --mode cache`. + +## Files + +- `c11_doppler_csi.py` shared library, time-varying TDL channel with + intra-symbol Doppler (genuine ICI), pilot aging, aging-aware LMMSE +- `c13_mnist.py` MNIST models (signed, Transformer SE, per-user AE), + training, MAML meta-training, SNR-sweep evaluation +- `c18_mnist_doppler.py` MNIST Doppler sweep +- `c19_mnist_epoch.py` convergence study with task-adapted checkpoints +- `c20_bert.py` BERT/AG News text study +- `c21_mnist_flat.py` flat Rayleigh study with decoder-side + first-order MAML +- `c22_maml_epoch.py` budget-matched meta-training trajectory +- `results_mnist/`, `results_bert/` stored CSV results and checkpoints +- `fig/` figure PDFs as used in the manuscript diff --git a/c11_doppler_csi.py b/c11_doppler_csi.py new file mode 100755 index 0000000..c43de88 --- /dev/null +++ b/c11_doppler_csi.py @@ -0,0 +1,804 @@ +#!/usr/bin/env python3 +# ============================================================ +# c11_doppler_csi.py +# +# TWC/TCOM revision experiments: +# Graceful degradation & CSI-aging robustness under a +# time-varying frequency-selective (TDL-like) OFDM channel. +# +# Frame structure (per frame): +# [ pilot OFDM symbol ] ... gap (DELTA-1 symbols) ... [ data OFDM symbol ] +# - Channel taps evolve continuously (Jakes sum-of-sinusoids) +# across the frame -> pilot CSI is OUTDATED at the data symbol. +# - Time-varying convolution within each OFDM symbol -> genuine ICI. +# +# Methods: +# qpsk_genie : OFDMA-QPSK, comb allocation (N/U subcarriers/user, +# repetition + MRC), perfect CSI at data time (bound) +# qpsk_pilot : same, but LS pilot CSI (aged) -> cliff effect +# joint : proposed user-wise attention semantic demux, +# joint training over (SNR, fD) grid, zero-shot eval +# joint_ft : joint + test-time fine-tuning (same budget as MAML) +# maml : SNR/Doppler-aware MAML (decoder-side inner loop) +# + test-time task-conditional adaptation +# +# The proposed receiver uses the SAME pilot (MMSE equalization with +# the aged LS estimate), so pilot overhead is identical to baseline. +# Difference is isolated to the demapping: fixed coherent QPSK vs +# learned semantic demultiplexing. +# +# Tasks tau = (SNR, fD_norm), fD_norm = f_D * T_sym (N samples). +# +# Usage: +# python3 c11_doppler_csi.py --mode train-joint +# python3 c11_doppler_csi.py --mode train-maml +# python3 c11_doppler_csi.py --mode eval +# python3 c11_doppler_csi.py --mode fig +# ============================================================ + +import argparse +import math +import os +import csv +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from torch.func import functional_call +except Exception: + from torch.nn.utils.stateless import functional_call + +try: + from scipy.special import j0 as bessel_j0 +except Exception: + def bessel_j0(x): + """Abramowitz & Stegun 9.4.1 / 9.4.3 polynomial approximation.""" + x = abs(float(x)) + if x < 3.0: + t = (x / 3.0) ** 2 + return (1.0 - 2.2499997 * t + 1.2656208 * t ** 2 - 0.3163866 * t ** 3 + + 0.0444479 * t ** 4 - 0.0039444 * t ** 5 + 0.0002100 * t ** 6) + t = 3.0 / x + f0 = (0.79788456 - 0.00000077 * t - 0.00552740 * t ** 2 - 0.00009512 * t ** 3 + + 0.00137237 * t ** 4 - 0.00072805 * t ** 5 + 0.00014476 * t ** 6) + th = (x - 0.78539816 - 0.04166397 * t - 0.00003954 * t ** 2 + 0.00262573 * t ** 3 + - 0.00054125 * t ** 4 - 0.00029333 * t ** 5 + 0.00013558 * t ** 6) + return f0 * math.cos(th) / math.sqrt(x) + + +# ------------------------------------------------------------ +# Reproducibility +# ------------------------------------------------------------ +def set_seed(seed: int = 0): + torch.manual_seed(seed) + np.random.seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +# ------------------------------------------------------------ +# Time-varying TDL channel (Jakes sum-of-sinusoids) +# ------------------------------------------------------------ +class TDLChannel: + """ + Exponential PDP, L taps, per-tap Jakes Doppler via sum-of-sinusoids. + fd_norm is the Doppler normalized to the useful OFDM symbol + duration (N samples): fd_norm = f_D * N * T_s. + """ + + def __init__(self, N=128, cp=16, L=8, pdp_decay=3.0, n_sin=16, device="cpu"): + self.N = N + self.cp = cp + self.L = L + self.S = N + cp + self.n_sin = n_sin + self.device = device + p = torch.exp(-torch.arange(L, dtype=torch.float32) / pdp_decay) + self.pdp = (p / p.sum()).to(device) # (L,) + + def sample(self, B, fd_norm, delta): + """ + Generate tap gains at pilot-symbol samples and data-symbol samples. + Pilot occupies samples [0, S); data occupies [delta*S, delta*S + S). + Returns g_p, g_d: (B, L, S) complex tap trajectories. + """ + dev = self.device + S, L, Ns = self.S, self.L, self.n_sin + fd_samp = fd_norm / self.N # per-sample normalized Doppler + + n_p = torch.arange(S, device=dev, dtype=torch.float32) + n_d = n_p + delta * S + n_all = torch.cat([n_p, n_d]) # (2S,) + + theta = 2 * math.pi * torch.rand(B, L, 1, Ns, device=dev) + phi = 2 * math.pi * torch.rand(B, L, 1, Ns, device=dev) + omega = 2 * math.pi * fd_samp * torch.cos(theta) # (B,L,1,Ns) + ph = omega * n_all.view(1, 1, -1, 1) + phi # (B,L,2S,Ns) + g = torch.exp(1j * ph).sum(dim=-1) / math.sqrt(Ns) # (B,L,2S) + g = g * torch.sqrt(self.pdp).view(1, L, 1) + return g[:, :, :S], g[:, :, S:] + + def transmit(self, X_freq, g, noise_var): + """ + One OFDM symbol through the time-varying channel. + X_freq: (B, N) complex subcarrier vector (unitary convention) + g: (B, L, S) tap gains over the symbol (incl. CP samples) + Returns Y: (B, N) complex received subcarrier vector. + """ + B = X_freq.shape[0] + N, cp, S, L = self.N, self.cp, self.S, self.L + + x = torch.fft.ifft(X_freq, dim=1) * math.sqrt(N) # (B,N) + x_cp = torch.cat([x[:, -cp:], x], dim=1) # (B,S) + + y = torch.zeros(B, S, dtype=torch.complex64, device=x.device) + for l in range(L): + if l == 0: + y = y + g[:, 0, :] * x_cp + else: + y[:, l:] = y[:, l:] + g[:, l, l:] * x_cp[:, :S - l] + + n = (torch.randn(B, S, device=x.device) + + 1j * torch.randn(B, S, device=x.device)) * math.sqrt(noise_var / 2.0) + y = y + n + + y_data = y[:, cp:] # discard CP + Y = torch.fft.fft(y_data, dim=1) / math.sqrt(N) # (B,N) + return Y + + def genie_H(self, g): + """Effective per-subcarrier channel: DFT of time-averaged taps.""" + g_bar = g[:, :, self.cp:].mean(dim=2) # (B,L) + h = torch.zeros(g.shape[0], self.N, dtype=torch.complex64, device=g.device) + h[:, :self.L] = g_bar + return torch.fft.fft(h, dim=1) # (B,N) + + +def make_pilot(N, device, seed=1234): + """Fixed pseudo-random QPSK pilot, |X_p[k]| = 1.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + idx = torch.randint(0, 4, (N,), generator=gen) + ang = math.pi / 4 + idx.to(torch.float32) * math.pi / 2 + return torch.exp(1j * ang).to(device) + + +# ------------------------------------------------------------ +# OFDMA-QPSK baseline (comb allocation + repetition + MRC) +# ------------------------------------------------------------ +QPSK = None # filled at runtime + + +def qpsk_constellation(device): + ang = math.pi / 4 + torch.arange(4, device=device, dtype=torch.float32) * math.pi / 2 + return torch.exp(1j * ang) # (4,) + + +def qpsk_tx(labels, U, N, device): + """labels: (B,U) in {0..3} -> X: (B,N), comb allocation.""" + B = labels.shape[0] + const = qpsk_constellation(device) + sym = const[labels] # (B,U) + X = torch.zeros(B, N, dtype=torch.complex64, device=device) + for u in range(U): + ks = torch.arange(u, N, U, device=device) + X[:, ks] = sym[:, u].unsqueeze(1) + return X + + +def qpsk_detect(Y, H_hat, U, N): + """MRC over each user's comb subcarriers with (possibly aged) CSI.""" + B = Y.shape[0] + const = qpsk_constellation(Y.device) # (4,) + preds = torch.zeros(B, U, dtype=torch.long, device=Y.device) + for u in range(U): + ks = torch.arange(u, N, U, device=Y.device) + Z = (torch.conj(H_hat[:, ks]) * Y[:, ks]).sum(dim=1) # (B,) + metric = torch.real(torch.conj(const).view(1, 4) * Z.unsqueeze(1)) + preds[:, u] = metric.argmax(dim=1) + return preds + + +# ------------------------------------------------------------ +# Proposed model: masked superposition + user-wise attention +# operating on the MMSE-equalized received subcarrier vector +# (real/imag stacked features, masks applied per subcarrier) +# ------------------------------------------------------------ +class UserWiseAttentionSem(nn.Module): + def __init__(self, U, V, d, hidden=256): + super().__init__() + self.U, self.V, self.d = U, V, d + self.codebook = nn.Parameter(torch.randn(V, d)) + self.masks = nn.Parameter(torch.randn(U, d)) + self.query = nn.Parameter(torch.randn(U, 2 * d)) + self.key = nn.Linear(2 * d, 2 * d, bias=False) + self.val = nn.Linear(2 * d, 2 * d, bias=False) + self.cls = nn.Sequential( + nn.Linear(2 * d, hidden), + nn.ReLU(inplace=True), + nn.Linear(hidden, V), + ) + + # ---- transmitter ---- + def tx(self, labels, params=None): + cb = self.codebook if params is None else params["codebook"] + mk = self.masks if params is None else params["masks"] + e = F.embedding(labels, cb) # (B,U,d) + m = F.normalize(mk, dim=1) # (U,d) + y = (e * m.unsqueeze(0)).sum(dim=1) # (B,d) + y = y / torch.sqrt(torch.mean(y ** 2, dim=1, keepdim=True) + 1e-12) + return y, m + + # ---- receiver ---- + def rx(self, Yeq, m, params=None): + B = Yeq.shape[0] + R = Yeq.unsqueeze(1) * m.unsqueeze(0) # (B,U,d) complex + phi = torch.cat([R.real, R.imag], dim=-1) # (B,U,2d) + + if params is None: + K = self.key(phi) + Vv = self.val(phi) + q = self.query + else: + K = F.linear(phi, params["key.weight"]) + Vv = F.linear(phi, params["val.weight"]) + q = params["query"] + + scores = torch.einsum("ud,bid->bui", q, K) / math.sqrt(2 * self.d) + attn = F.softmax(scores, dim=-1) # (B,U,U) + z = torch.einsum("bui,bid->bud", attn, Vv) + phi # residual (B,U,2d) + + if params is None: + logits = self.cls(z) + else: + h1 = F.linear(z, params["cls.0.weight"], params["cls.0.bias"]) + h1 = F.relu(h1) + logits = F.linear(h1, params["cls.2.weight"], params["cls.2.bias"]) + return logits # (B,U,V) + + +class TransformerSem(UserWiseAttentionSem): + """ + SOTA baseline: shared-embedding TX (same masking) + Transformer + encoder separation treating users as tokens (JSAC'26 [21]-style). + """ + + def __init__(self, U, V, d, hidden=256, n_heads=4, n_layers=2): + super().__init__(U, V, d, hidden) + self.inp = nn.Linear(2 * d, hidden) + layer = nn.TransformerEncoderLayer(d_model=hidden, nhead=n_heads, + dim_feedforward=2 * hidden, + batch_first=True) + self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers) + self.out = nn.Linear(hidden, V) + + def rx(self, Yeq, m, params=None): + R = Yeq.unsqueeze(1) * m.unsqueeze(0) # (B,U,d) + phi = torch.cat([R.real, R.imag], dim=-1) # (B,U,2d) + z = self.encoder(self.inp(phi)) # (B,U,hidden) + return self.out(z) # (B,U,V) + + +class PerUserAESem(nn.Module): + """ + SOTA baseline: DeepMA-style per-user autoencoder multiple access. + Each user has a dedicated codebook (no shared masking) and a + dedicated decoder head operating on the equalized observation. + """ + + def __init__(self, U, V, d, hidden=256): + super().__init__() + self.U, self.V, self.d = U, V, d + self.codebooks = nn.Parameter(torch.randn(U, V, d)) + self.heads = nn.ModuleList([ + nn.Sequential(nn.Linear(2 * d, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, V)) + for _ in range(U) + ]) + + def tx(self, labels, params=None): + B = labels.shape[0] + idx = labels + torch.arange(self.U, device=labels.device).view(1, -1) * self.V + e = self.codebooks.view(self.U * self.V, self.d)[idx] # (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) # (B,U,V) + + +class SignedAttentionSem(UserWiseAttentionSem): + """ + Proposed signed user-wise attention. The candidate Gram matrix + T = Phi Phi^T / (2d) is mapped by a small score network to signed + combining weights W = I + f_theta(T), which can subtract correlated + interference, unlike convex nonnegative softmax weights. The + query/key/value projections of the parent are unused. + """ + + def __init__(self, U, V, d, hidden=256, score_hidden=64): + super().__init__(U, V, d, hidden) + 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), + ) + + def rx(self, Yeq, m, params=None): + B = Yeq.shape[0] + R = Yeq.unsqueeze(1) * m.unsqueeze(0) # (B,U,d) + phi = torch.cat([R.real, R.imag], dim=-1) # (B,U,2d) + T = torch.bmm(phi, phi.transpose(1, 2)) / (2 * self.d) + + if params is None: + 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) + + h = T.reshape(B, -1) + h = F.relu(F.linear(h, params["score.0.weight"], params["score.0.bias"])) + h = F.relu(F.linear(h, params["score.2.weight"], params["score.2.bias"])) + w = F.linear(h, params["score.4.weight"], params["score.4.bias"]) + W = torch.eye(self.U, device=phi.device).unsqueeze(0) + w.view(B, self.U, self.U) + z = torch.bmm(W, phi) + h1 = F.relu(F.linear(z, params["cls.0.weight"], params["cls.0.bias"])) + return F.linear(h1, params["cls.2.weight"], params["cls.2.bias"]) + + +MODEL_CLASSES = { + "attention": UserWiseAttentionSem, + "signed": SignedAttentionSem, + "transformer": TransformerSem, + "peruser": PerUserAESem, +} + +DECODER_KEYS = ["query", "key.weight", "val.weight", + "cls.0.weight", "cls.0.bias", "cls.2.weight", "cls.2.bias"] + +SIGNED_DECODER_KEYS = ["score.0.weight", "score.0.bias", + "score.2.weight", "score.2.bias", + "score.4.weight", "score.4.bias", + "cls.0.weight", "cls.0.bias", + "cls.2.weight", "cls.2.bias"] + + +def decoder_keys_for(model): + return SIGNED_DECODER_KEYS if isinstance(model, SignedAttentionSem) else DECODER_KEYS + + +def get_params(model, keys=None): + d = dict(model.named_parameters()) + if keys is None: + return d + return {k: d[k] for k in keys} + + +# ------------------------------------------------------------ +# End-to-end forward through channel (differentiable) +# ------------------------------------------------------------ +def aging_rho(fd_norm, delta, N, cp): + """Pilot-to-data temporal correlation: rho = J0(2*pi*fd_samp*lag).""" + fd_samp = fd_norm / N + lag = delta * (N + cp) + return float(bessel_j0(2.0 * math.pi * fd_samp * lag)) + + +def semantic_forward(model, labels, chan, g_p, g_d, X_pilot, noise_var, fd_norm, delta, + params=None): + """ + Full pipeline: TX -> time-varying channel (pilot + data symbols) + -> LS pilot estimate -> aging-aware LMMSE equalization + -> user-wise attention. + + Aging-aware LMMSE: with H_d = rho*H_p + sqrt(1-rho^2)*innovation and + LS estimate Hhat = H_p + e (Var e = sigma^2), the LMMSE predictor of + the data-time channel is Htil = (rho/(1+sigma^2))*Hhat with residual + variance q = 1 - rho^2/(1+sigma^2). The equalizer stays regularized + at high SNR because q > 0 whenever rho < 1. + """ + full = None + if params is not None: + full = dict(model.named_parameters()) + full.update(params) + + y_emb, m = model.tx(labels, params=full) + X_data = y_emb.to(torch.complex64) # real coeffs, imag=0 + + Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(labels.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) # |X_p|=1 + + rho = aging_rho(fd_norm, 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) + + # per-frame RMS normalization: removes the task-dependent scale of the + # LMMSE output (~rho/(1+sigma^2)) so that a single receiver operates on + # a scale-invariant representation across (SNR, Doppler) tasks + rms = torch.sqrt(torch.mean(Yeq.abs() ** 2, dim=1, keepdim=True) + 1e-12) + Yeq = Yeq / rms + + return model.rx(Yeq, m, params=full) + + +# ------------------------------------------------------------ +# Task sampling +# ------------------------------------------------------------ +def sample_task(args, rng): + snr = float(rng.choice(args.train_snrs)) + fd = float(rng.choice(args.train_fds)) + return snr, fd + + +# ------------------------------------------------------------ +# Training +# ------------------------------------------------------------ +def run_batch_loss(model, chan, X_pilot, args, snr, fd, batch, params=None, device="cpu"): + labels = torch.randint(0, args.vocab, (batch, args.users), device=device) + g_p, g_d = chan.sample(batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = semantic_forward(model, labels, chan, g_p, g_d, X_pilot, noise_var, + fd, args.delta, params=params) + loss = F.cross_entropy(logits.reshape(-1, args.vocab), labels.reshape(-1)) + return loss, logits, labels + + +def train_joint(args, device, model_key="attention", ckpt="joint.pt"): + set_seed(args.seed) + chan = TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = make_pilot(args.nfft, device) + model = MODEL_CLASSES[model_key](args.users, args.vocab, 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 = sample_task(args, rng) + loss, _, _ = run_batch_loss(model, chan, X_pilot, args, snr, fd, args.batch, device=device) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + if step % 200 == 0: + print(f"[{model_key} {step}/{args.steps}] loss={loss.item():.4f} (snr={snr}, fd={fd})", flush=True) + + os.makedirs(args.save_dir, exist_ok=True) + torch.save(model.state_dict(), os.path.join(args.save_dir, ckpt)) + print(f"saved {ckpt}") + + +def train_maml(args, device, model_key="attention", ckpt="maml.pt"): + set_seed(args.seed) + chan = TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = make_pilot(args.nfft, device) + model = MODEL_CLASSES[model_key](args.users, args.vocab, args.dim, args.hidden).to(device) + opt = torch.optim.Adam(model.parameters(), lr=args.meta_lr) + rng = np.random.default_rng(args.seed) + dkeys = decoder_keys_for(model) + + for step in range(1, args.meta_steps + 1): + meta_loss = 0.0 + for _ in range(args.meta_batch): + snr, fd = sample_task(args, rng) + fast = get_params(model, dkeys) + for _ in range(args.inner_steps): + loss_sup, _, _ = run_batch_loss(model, chan, X_pilot, args, snr, fd, + args.batch, params=fast, device=device) + grads = torch.autograd.grad(loss_sup, list(fast.values()), create_graph=False) + fast = {k: p - args.inner_lr * g.detach() + for (k, p), g in zip(fast.items(), grads)} + loss_q, _, _ = run_batch_loss(model, chan, X_pilot, args, snr, fd, + args.batch, params=fast, device=device) + meta_loss = meta_loss + loss_q + meta_loss = meta_loss / args.meta_batch + + opt.zero_grad(set_to_none=True) + meta_loss.backward() + opt.step() + if step % 100 == 0: + print(f"[maml {step}/{args.meta_steps}] meta-loss={meta_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(f"saved {ckpt}") + + +# ------------------------------------------------------------ +# Evaluation +# ------------------------------------------------------------ +def adapt(model, chan, X_pilot, args, snr, fd, device): + """Test-time task-conditional adaptation (decoder-side, FOMAML-style).""" + fast = {k: v.detach().clone() for k, v in get_params(model, decoder_keys_for(model)).items()} + for k in fast: + fast[k].requires_grad_(True) + for _ in range(args.eval_inner_steps): + loss, _, _ = run_batch_loss(model, chan, X_pilot, args, snr, fd, + args.support, params=fast, device=device) + 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()} + + +@torch.no_grad() +def eval_semantic(model, chan, X_pilot, args, snr, fd, n_batches, device, params=None): + errs, total = 0, 0 + for _ in range(n_batches): + labels = torch.randint(0, args.vocab, (args.eval_batch, args.users), device=device) + g_p, g_d = chan.sample(args.eval_batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = semantic_forward(model, labels, chan, g_p, g_d, X_pilot, noise_var, + fd, args.delta, params=params) + preds = logits.argmax(dim=-1) + errs += (preds != labels).sum().item() + total += labels.numel() + return errs / total + + +@torch.no_grad() +def eval_qpsk(chan, X_pilot, args, snr, fd, n_batches, device, genie=False): + errs, total = 0, 0 + for _ in range(n_batches): + labels = torch.randint(0, args.vocab, (args.eval_batch, args.users), device=device) + g_p, g_d = chan.sample(args.eval_batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + + X_d = qpsk_tx(labels, args.users, args.nfft, device) + Y_d = chan.transmit(X_d, g_d, noise_var) + + if genie: + H_hat = chan.genie_H(g_d) + else: + Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(labels.shape[0], -1), g_p, noise_var) + H_hat = Y_p * torch.conj(X_pilot).unsqueeze(0) + + preds = qpsk_detect(Y_d, H_hat, args.users, args.nfft) + errs += (preds != labels).sum().item() + total += labels.numel() + return errs / total + + +def eval_all(args, device): + set_seed(args.seed + 1) + chan = TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = make_pilot(args.nfft, device) + + joint = UserWiseAttentionSem(args.users, args.vocab, args.dim, args.hidden).to(device) + joint.load_state_dict(torch.load(os.path.join(args.save_dir, "joint.pt"), map_location=device)) + joint.eval() + maml = UserWiseAttentionSem(args.users, args.vocab, args.dim, args.hidden).to(device) + maml.load_state_dict(torch.load(os.path.join(args.save_dir, "maml.pt"), map_location=device)) + maml.eval() + + signed_models = {} + for key, ckpt in [("sjoint", "signed_joint.pt"), ("smaml", "signed_maml.pt")]: + path = os.path.join(args.save_dir, ckpt) + if os.path.exists(path): + mdl = SignedAttentionSem(args.users, args.vocab, args.dim, args.hidden).to(device) + mdl.load_state_dict(torch.load(path, map_location=device)) + mdl.eval() + signed_models[key] = mdl + + # SOTA comparison models (zero-shot joint training), if available + sota = {} + for key, ckpt in [("transformer", "transformer.pt"), ("peruser", "peruser.pt")]: + path = os.path.join(args.save_dir, ckpt) + if os.path.exists(path): + mdl = MODEL_CLASSES[key](args.users, args.vocab, args.dim, args.hidden).to(device) + mdl.load_state_dict(torch.load(path, map_location=device)) + mdl.eval() + sota[key] = mdl + + sweeps = [] + # Sweep A: SER vs Doppler at fixed SNR + for fd in args.eval_fds: + sweeps.append(("doppler", args.eval_snr_fixed, fd)) + # Sweep B: SER vs SNR at fixed (high) Doppler + for snr in args.eval_snrs: + sweeps.append(("snr", snr, args.eval_fd_fixed)) + + os.makedirs(args.save_dir, exist_ok=True) + csv_path = os.path.join(args.save_dir, "doppler_csi_results.csv") + with open(csv_path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["sweep", "snr_db", "fd_norm", "method", "ser"]) + + for sweep, snr, fd in sweeps: + row = {} + row["qpsk_genie"] = eval_qpsk(chan, X_pilot, args, snr, fd, args.eval_nb, device, genie=True) + row["qpsk_pilot"] = eval_qpsk(chan, X_pilot, args, snr, fd, args.eval_nb, device, genie=False) + for key, mdl in sota.items(): + row[key] = eval_semantic(mdl, chan, X_pilot, args, snr, fd, args.eval_nb, device) + row["joint"] = eval_semantic(joint, chan, X_pilot, args, snr, fd, args.eval_nb, device) + if "sjoint" in signed_models: + row["sjoint"] = eval_semantic(signed_models["sjoint"], chan, X_pilot, + args, snr, fd, args.eval_nb, device) + + # adaptation-based methods: average over independent trials + adapt_list = [("joint_ft", joint), ("maml", maml)] + if "smaml" in signed_models: + adapt_list.append(("smaml", signed_models["smaml"])) + for name, mdl in adapt_list: + sers = [] + for _ in range(args.adapt_trials): + with torch.enable_grad(): + fast = adapt(mdl, chan, X_pilot, args, snr, fd, device) + sers.append(eval_semantic(mdl, chan, X_pilot, args, snr, fd, + args.eval_nb_adapt, device, params=fast)) + row[name] = float(np.mean(sers)) + + for method, ser in row.items(): + w.writerow([sweep, snr, fd, method, ser]) + f.flush() + print(f"[{sweep}] snr={snr:5.1f} fd={fd:6.3f} | " + + " ".join(f"{k}={v:.4e}" for k, v in row.items()), flush=True) + + print(f"saved {csv_path}") + + +# ------------------------------------------------------------ +# Figures +# ------------------------------------------------------------ +LABELS = { + "qpsk_genie": "OFDMA-QPSK genie CSI", + "qpsk_pilot": "OFDMA-QPSK pilot CSI", + "transformer": "Transformer SE separation", + "peruser": "Per-user AE multiple access", + "joint": "Proposed softmax joint", + "joint_ft": "Proposed softmax fine-tuned", + "maml": "Proposed softmax MAML", + "sjoint": "Proposed signed joint", + "smaml": "Proposed signed MAML", +} +STYLES = { + "qpsk_genie": dict(color="gray", marker="^", ls="--"), + "qpsk_pilot": dict(color="k", marker="v", ls="-"), + "transformer": dict(color="tab:purple", marker="P", ls="-"), + "peruser": dict(color="tab:brown", marker="X", ls="-"), + "joint": dict(color="tab:blue", marker="s", ls="-"), + "joint_ft": dict(color="tab:green", marker="D", ls="-"), + "maml": dict(color="tab:cyan", marker="d", ls="-"), + "sjoint": dict(color="tab:red", marker="o", ls="-"), + "smaml": dict(color="tab:green", marker="D", ls="-"), +} +PLOT_KEYS = ["qpsk_genie", "qpsk_pilot", "transformer", "peruser", + "joint", "sjoint", "smaml"] + + +def make_figs(args): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + csv_path = os.path.join(args.save_dir, "doppler_csi_results.csv") + rows = [] + with open(csv_path) as f: + for r in csv.DictReader(f): + rows.append(r) + + os.makedirs(args.fig_dir, exist_ok=True) + + # Fig A: SER vs normalized Doppler + fig = plt.figure(figsize=(5.2, 3.9)) + ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) + for m in PLOT_KEYS: + latest = {} + for r in rows: + if r["sweep"] == "doppler" and r["method"] == m: + latest[float(r["fd_norm"])] = float(r["ser"]) + pts = sorted((k, v) for k, v in latest.items() if v > 0) + if pts: + x, y = zip(*pts) + ax.semilogy(x, y, label=LABELS[m], ms=4, lw=1.2, **STYLES[m]) + ax.set_xscale("log") + 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, loc="upper left") + fig.savefig(os.path.join(args.fig_dir, f"ser_vs_doppler_snr{int(args.eval_snr_fixed)}.pdf")) + print("saved fig A") + + # Fig B: SER vs SNR at fixed high Doppler + fig = plt.figure(figsize=(5.2, 3.9)) + ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) + for m in PLOT_KEYS: + latest = {} + for r in rows: + if r["sweep"] == "snr" and r["method"] == m: + latest[float(r["snr_db"])] = float(r["ser"]) + pts = sorted((k, v) for k, v in latest.items() if v > 0) + if pts: + x, y = zip(*pts) + ax.semilogy(x, y, label=LABELS[m], ms=4, lw=1.2, **STYLES[m]) + ax.set_xlabel("SNR (dB)") + ax.set_ylabel("SER") + ax.grid(True, which="both", alpha=0.35) + ax.legend(fontsize=7, loc="lower center") + fig.savefig(os.path.join(args.fig_dir, f"ser_vs_snr_fd{args.eval_fd_fixed}.pdf")) + print("saved fig B") + + +# ------------------------------------------------------------ +# Main +# ------------------------------------------------------------ +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["train-joint", "train-maml", "train-transformer", + "train-peruser", "train-signed-joint", + "train-signed-maml", "eval", "fig"], required=True) + + # system + p.add_argument("--users", type=int, default=8) + p.add_argument("--vocab", type=int, default=4) + 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, help="pilot-to-data gap (OFDM symbols)") + + # task grids + 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]) + + # training + p.add_argument("--steps", type=int, default=6000) + p.add_argument("--batch", type=int, default=256) + p.add_argument("--lr", type=float, default=1e-3) + + # maml + p.add_argument("--meta-steps", type=int, default=3000) + p.add_argument("--meta-batch", type=int, default=4) + p.add_argument("--inner-steps", type=int, default=1) + p.add_argument("--inner-lr", type=float, default=0.05) + p.add_argument("--meta-lr", type=float, default=1e-3) + + # evaluation + p.add_argument("--eval-batch", type=int, default=256) + p.add_argument("--eval-nb", type=int, default=200, help="eval batches (zero-shot)") + p.add_argument("--eval-nb-adapt", type=int, default=40, help="eval batches per adapt trial") + p.add_argument("--adapt-trials", type=int, default=8) + p.add_argument("--support", type=int, default=64, help="support frames per adapt step") + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.02) + + p.add_argument("--eval-fds", type=float, nargs="+", + default=[0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1]) + p.add_argument("--eval-snr-fixed", type=float, default=15.0) + p.add_argument("--eval-snrs", type=float, nargs="+", default=[0, 5, 10, 15, 20, 25, 30]) + p.add_argument("--eval-fd-fixed", type=float, default=0.05) + + # misc + p.add_argument("--save-dir", type=str, default="results_doppler") + p.add_argument("--fig-dir", type=str, default="fig") + p.add_argument("--seed", type=int, default=0) + args = p.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"device: {device}") + + if args.mode == "train-joint": + train_joint(args, device) + elif args.mode == "train-transformer": + train_joint(args, device, model_key="transformer", ckpt="transformer.pt") + elif args.mode == "train-peruser": + train_joint(args, device, model_key="peruser", ckpt="peruser.pt") + elif args.mode == "train-signed-joint": + train_joint(args, device, model_key="signed", ckpt="signed_joint.pt") + elif args.mode == "train-signed-maml": + train_maml(args, device, model_key="signed", ckpt="signed_maml.pt") + elif args.mode == "train-maml": + train_maml(args, device) + elif args.mode == "eval": + eval_all(args, device) + elif args.mode == "fig": + make_figs(args) + + +if __name__ == "__main__": + main() diff --git a/c13_mnist.py b/c13_mnist.py new file mode 100755 index 0000000..1200406 --- /dev/null +++ b/c13_mnist.py @@ -0,0 +1,586 @@ +#!/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() diff --git a/c18_mnist_doppler.py b/c18_mnist_doppler.py new file mode 100755 index 0000000..3108d6e --- /dev/null +++ b/c18_mnist_doppler.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +# ============================================================ +# c18_mnist_doppler.py +# +# MNIST real-data SER versus normalized Doppler at a fixed SNR, +# reusing the trained checkpoints of c13_mnist.py. The signed +# MAML receiver re-adapts its decoder-side parameters for every +# (SNR, Doppler) task. +# ============================================================ + +import argparse +import csv +import os +import numpy as np +import torch + +import c11_doppler_csi as base +import c13_mnist as m13 + + +@torch.no_grad() +def run_sweep(args, device): + base.set_seed(args.seed + 3) + _, te = m13.get_datasets(args.data_root) + x, y = m13.tensorize(te) + xs_tr, ys_tr = m13.tensorize(m13.get_datasets(args.data_root)[0]) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + + sem = m13.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() + mm = m13.MnistSemanticMA(args.users, args.dim, args.hidden).to(device) + mm.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_maml.pt"), + map_location=device)) + mm.eval() + tfm = m13.MnistTransformerMA(args.users, args.dim, args.hidden).to(device) + tfm.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_tf.pt"), + map_location=device)) + tfm.eval() + aem = m13.MnistPerUserAE(args.users, args.dim, args.hidden).to(device) + aem.load_state_dict(torch.load(os.path.join(args.save_dir, "mnist_ae.pt"), + map_location=device)) + aem.eval() + txc = m13.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() + + snr = args.eval_snr + noise_var = 10 ** (-snr / 10.0) + rng = np.random.default_rng(args.seed + 3) + csv_path = os.path.join(args.save_dir, "mnist_doppler.csv") + with open(csv_path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["snr_db", "fd_norm", "method", "ser"]) + for fd in args.eval_fds: + args.eval_fd = fd + errs = {"digital_genie": 0, "digital_pilot": 0, "semantic": 0, + "semantic_tf": 0, "semantic_ae": 0, "semantic_maml": 0} + fast_mm = m13.adapt_mnist(mm, xs_tr, ys_tr, chan, X_pilot, + args, snr, device) + total = 0 + for _ in range(args.eval_nb): + imgs, labels = m13.sample_frames(x, y, args.eval_batch, + args.users, device, rng) + g_p, g_d = chan.sample(args.eval_batch, fd, args.delta) + + logits = m13.forward_frames(sem, imgs, chan, g_p, g_d, + X_pilot, noise_var, fd, args.delta) + errs["semantic"] += (logits.argmax(-1) != labels).sum().item() + lg2 = m13.forward_frames(tfm, imgs, chan, g_p, g_d, X_pilot, + noise_var, fd, args.delta) + errs["semantic_tf"] += (lg2.argmax(-1) != labels).sum().item() + lg4 = m13.forward_frames(aem, imgs, chan, g_p, g_d, X_pilot, + noise_var, fd, args.delta) + errs["semantic_ae"] += (lg4.argmax(-1) != labels).sum().item() + lg3 = m13.forward_frames_fast(mm, fast_mm, imgs, chan, g_p, + g_d, X_pilot, noise_var, fd, + args.delta) + errs["semantic_maml"] += (lg3.argmax(-1) != labels).sum().item() + + 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 = m13.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 = m13.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, fd, name, e / total]) + f.flush() + print(f"fd={fd:7.4f} | " + + " ".join(f"{k}={v / total:.4e}" for k, v in errs.items()), + flush=True) + print("saved", csv_path) + + +def make_fig(args): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + rows = list(csv.DictReader(open(os.path.join(args.save_dir, + "mnist_doppler.csv")))) + 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["fd_norm"]), float(r["ser"])) + for r in rows if r["method"] == mkey]) + pts = [(a, b) for a, b in pts if b > 0] + if not pts: + continue + xs, ys = zip(*pts) + ax.semilogy(xs, ys, label=LAB[mkey], ms=4, lw=1.3, **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)) + out = os.path.join(args.fig_dir, + f"mnist_ser_vs_doppler_snr{int(args.eval_snr)}.pdf") + fig.savefig(out) + print("saved", out) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["run", "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("--eval-snr", type=float, default=15.0) + p.add_argument("--eval-fds", type=float, nargs="+", + default=[0.001, 0.002, 0.005, 0.01, 0.02, 0.03, 0.05, + 0.0567, 0.07, 0.085, 0.1]) + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.01) + p.add_argument("--support", type=int, default=32) + 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() + args.eval_fd = args.eval_fds[0] + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("device:", device) + if args.mode == "run": + run_sweep(args, device) + else: + make_fig(args) + + +if __name__ == "__main__": + main() diff --git a/c19_mnist_epoch.py b/c19_mnist_epoch.py new file mode 100755 index 0000000..2586c23 --- /dev/null +++ b/c19_mnist_epoch.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +# ============================================================ +# c19_mnist_epoch.py +# +# Real-data convergence study: evaluation SER versus training +# step for the proposed signed receiver, the Transformer SE +# separator, and the per-user AE, all trained on the MNIST +# TDL chain with the identical protocol. The digital chain +# floors (genie / aged pilot) are drawn as horizontal +# references since they involve no training. +# ============================================================ + +import argparse +import csv +import os +import numpy as np +import torch + +import c11_doppler_csi as base +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"), +} + + +@torch.no_grad() +def eval_ser(model, x, y, chan, X_pilot, args, device, rng): + model.eval() + noise_var = 10 ** (-args.eval_snr / 10.0) + errs, total = 0, 0 + for _ in range(args.eval_nb): + imgs, labels = m13.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) + logits = m13.forward_frames(model, imgs, chan, g_p, g_d, X_pilot, + noise_var, args.eval_fd, args.delta) + errs += (logits.argmax(-1) != labels).sum().item() + total += labels.numel() + model.train() + return errs / total + + +@torch.no_grad() +def eval_ser_adapted(model, fast, x, y, chan, X_pilot, args, device): + noise_var = 10 ** (-args.eval_snr / 10.0) + rng = np.random.default_rng(args.seed + 3) + errs, total = 0, 0 + for _ in range(args.eval_nb): + imgs, labels = m13.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) + logits = m13.forward_frames_fast(model, fast, imgs, chan, g_p, g_d, + X_pilot, noise_var, args.eval_fd, + args.delta) + errs += (logits.argmax(-1) != labels).sum().item() + total += labels.numel() + return errs / total + + +def run_maml_track(args, device): + """Retrace the signed joint training and record the task-adapted + SER at every checkpoint, appending method 'semantic_maml_pre'.""" + tr, te = m13.get_datasets(args.data_root) + x_tr, y_tr = m13.tensorize(tr) + x_te, y_te = m13.tensorize(te) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + base.set_seed(args.seed) + model = m13.MnistSemanticMA(args.users, args.dim, args.hidden).to(device) + opt = torch.optim.Adam(model.parameters(), lr=args.lr) + rng = np.random.default_rng(args.seed) + csv_path = os.path.join(args.save_dir, "mnist_epoch.csv") + with open(csv_path, "a", newline="") as f: + w = csv.writer(f) + fast = m13.adapt_mnist(model, x_tr, y_tr, chan, X_pilot, args, + args.eval_snr, device) + ser0 = eval_ser_adapted(model, fast, x_te, y_te, chan, X_pilot, + args, device) + w.writerow(["semantic_maml_pre", 0, ser0]) + for step in range(1, args.steps + 1): + snr = float(rng.choice(args.train_snrs)) + fd = float(rng.choice(args.train_fds)) + imgs, labels = m13.sample_frames(x_tr, y_tr, args.batch, + args.users, device, rng) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = m13.forward_frames(model, imgs, chan, g_p, g_d, + X_pilot, noise_var, fd, args.delta) + loss = torch.nn.functional.cross_entropy( + logits.reshape(-1, 10), labels.reshape(-1)) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + if step % args.ckpt_every == 0: + fast = m13.adapt_mnist(model, x_tr, y_tr, chan, X_pilot, + args, args.eval_snr, device) + ser = eval_ser_adapted(model, fast, x_te, y_te, chan, + X_pilot, args, device) + w.writerow(["semantic_maml_pre", step, ser]) + f.flush() + print(f"[maml-pre {step}/{args.steps}] ser={ser:.4e}", + flush=True) + print("appended semantic_maml_pre to", csv_path) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["run", "run-maml", "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=10000) + p.add_argument("--ckpt-every", type=int, default=500) + p.add_argument("--batch", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--eval-snr", type=float, default=15.0) + p.add_argument("--eval-fd", type=float, default=0.05) + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.01) + p.add_argument("--support", type=int, default=32) + p.add_argument("--eval-batch", type=int, default=64) + p.add_argument("--eval-nb", type=int, default=30) + 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) + + csv_path = os.path.join(args.save_dir, "mnist_epoch.csv") + + if args.mode == "run-maml": + run_maml_track(args, device) + return + + if args.mode == "run": + tr, te = m13.get_datasets(args.data_root) + x_tr, y_tr = m13.tensorize(tr) + x_te, y_te = m13.tensorize(te) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + with open(csv_path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["method", "step", "ser"]) + for key, (cls, _) in MODELS.items(): + base.set_seed(args.seed) + 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) + erng = np.random.default_rng(args.seed + 3) + ser0 = eval_ser(model, x_te, y_te, chan, X_pilot, args, + device, np.random.default_rng(args.seed + 3)) + w.writerow([key, 0, ser0]) + for step in range(1, args.steps + 1): + snr = float(rng.choice(args.train_snrs)) + fd = float(rng.choice(args.train_fds)) + imgs, labels = m13.sample_frames(x_tr, y_tr, args.batch, + args.users, device, rng) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = m13.forward_frames(model, imgs, chan, g_p, g_d, + X_pilot, noise_var, fd, + args.delta) + loss = torch.nn.functional.cross_entropy( + logits.reshape(-1, 10), labels.reshape(-1)) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + if step % args.ckpt_every == 0: + ser = eval_ser(model, x_te, y_te, chan, X_pilot, + args, device, + np.random.default_rng(args.seed + 3)) + w.writerow([key, step, ser]) + f.flush() + print(f"[{key} {step}/{args.steps}] ser={ser:.4e}", + flush=True) + print("saved", csv_path) + return + + # ---- fig ---- + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + rows = list(csv.DictReader(open(csv_path))) + dig = list(csv.DictReader(open(os.path.join(args.save_dir, + "mnist_results.csv")))) + floors = {r["method"]: float(r["ser"]) for r in dig + if float(r["snr_db"]) == args.eval_snr + and r["method"] in ("digital_genie", "digital_pilot")} + 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="-")} + fig = plt.figure(figsize=(5.2, 3.9)) + ax = fig.add_axes([0.14, 0.125, 0.835, 0.845]) + 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, + 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, + 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]) + # 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 + # shown, and right of the dotted line decoder-side meta-training + # continues within the same total step budget. + maml_path = os.path.join(args.save_dir, "mnist_maml_epoch.csv") + if os.path.exists(maml_path): + mrows = list(csv.DictReader(open(maml_path))) + meta_pts = sorted([(int(r["step"]), float(r["ser"])) + for r in mrows]) + warm_end = meta_pts[0][0] + pre_pts = sorted([(int(r["step"]), float(r["ser"])) + for r in rows + if r["method"] == "semantic_maml_pre" + and int(r["step"]) < warm_end]) + if not pre_pts: + 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) + ax.set_xlabel("Training step") + ax.set_ylabel("SER") + if "digital_genie" in floors: + 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)) + out = os.path.join(args.fig_dir, "mnist_epoch_convergence.pdf") + fig.savefig(out) + print("saved", out) + + +if __name__ == "__main__": + main() diff --git a/c20_bert.py b/c20_bert.py new file mode 100755 index 0000000..fe18a37 --- /dev/null +++ b/c20_bert.py @@ -0,0 +1,533 @@ +#!/usr/bin/env python3 +# ============================================================ +# c20_bert.py +# +# Real-text semantic transmission with BERT features over the +# time-varying TDL channel. Each of the U users transmits one +# AG News sentence: a frozen BERT-base encoder produces the +# [CLS] feature (768-dim), a shared trainable projection maps +# it to the d=128 embedding, and the masked embeddings are +# superposed exactly as in the MNIST study. The receiver +# recovers each user's news topic (4 classes = 2 bits), and +# the conventional digital chain classifies at the +# transmitter and sends the 2-bit class index as one QPSK +# symbol over the user's comb subcarriers with 16-fold +# repetition and MRC. +# +# Modes: cache -> train / train-tf / train-ae / train-tx-cls +# -> train-maml -> eval -> fig +# ============================================================ + +import argparse +import csv +import os +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +import c11_doppler_csi as base +import c13_mnist as m13 + + +# ------------------------------------------------------------ +# BERT feature cache +# ------------------------------------------------------------ +def cache_features(args, device): + from transformers import AutoTokenizer, AutoModel + from datasets import load_dataset + tok = AutoTokenizer.from_pretrained("bert-base-uncased") + bert = AutoModel.from_pretrained("bert-base-uncased").to(device).eval() + ds = load_dataset("fancyzhx/ag_news") + out = {} + for split, n in [("train", args.cache_train), ("test", args.cache_test)]: + texts = ds[split]["text"][:n] + labels = torch.tensor(ds[split]["label"][:n]) + feats = [] + with torch.no_grad(): + for i in range(0, len(texts), 128): + bt = tok(texts[i:i + 128], padding=True, truncation=True, + max_length=64, return_tensors="pt").to(device) + cls = bert(**bt).last_hidden_state[:, 0] + feats.append(cls.cpu()) + if (i // 128) % 20 == 0: + print(f"[cache {split}] {i}/{len(texts)}", flush=True) + out[split] = (torch.cat(feats), labels) + os.makedirs(args.save_dir, exist_ok=True) + torch.save(out, os.path.join(args.save_dir, "bert_feats.pt")) + print("saved bert_feats.pt", + out["train"][0].shape, out["test"][0].shape) + + +def load_features(args): + d = torch.load(os.path.join(args.save_dir, "bert_feats.pt"), + map_location="cpu") + return d["train"], d["test"] + + +def sample_frames(feats, labels, B, U, device, rng): + idx = torch.from_numpy(rng.integers(0, feats.shape[0], size=(B * U,))) + x = feats[idx].to(device).view(B, U, -1) + y = labels[idx].to(device).view(B, U) + return x, y + + +# ------------------------------------------------------------ +# Models (mirror c13 with a BERT-feature front end) +# ------------------------------------------------------------ +class BertTrunk(nn.Module): + def __init__(self, out_dim, in_dim=768, hidden=256): + super().__init__() + self.net = nn.Sequential( + nn.Linear(in_dim, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, out_dim)) + + def forward(self, x): + return self.net(x) + + +class BertSemanticMA(nn.Module): + """Shared projection + masks + signed user-wise attention.""" + + def __init__(self, U, d=128, hidden=256, n_cls=4, score_hidden=64): + super().__init__() + self.U, self.d = U, d + self.encoder = BertTrunk(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, x, params=None): + B = x.shape[0] + e = self.encoder(x.reshape(B * self.U, -1)).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) + + +class BertTransformerMA(BertSemanticMA): + def __init__(self, U, d=128, hidden=256, n_cls=4, n_heads=4, + n_layers=2): + super().__init__(U, d, hidden, n_cls) + 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): + 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 BertPerUserAE(nn.Module): + """Per-user projections and heads, no masks.""" + + def __init__(self, U, d=128, hidden=256, n_cls=4): + super().__init__() + self.U, self.d = U, d + self.encs = nn.ModuleList([BertTrunk(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, x, params=None): + e = torch.stack([self.encs[u](x[:, u]) for u in range(self.U)], + dim=1) + 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) + return torch.stack([h(phi) for h in self.heads], dim=1) + + +class BertTxClassifier(nn.Module): + def __init__(self, n_cls=4, hidden=256): + super().__init__() + self.net = nn.Sequential( + nn.Linear(768, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, n_cls)) + + def forward(self, x): + return self.net(x) + + +# ------------------------------------------------------------ +# Digital chain: 2-bit class as one QPSK symbol, 16x repetition +# ------------------------------------------------------------ +def digital_tx(cls_idx, U, N, device): + const = torch.tensor([1 + 1j, 1 - 1j, -1 + 1j, -1 - 1j], + device=device) / np.sqrt(2.0) + B = cls_idx.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) + X[:, ks] = const[cls_idx[:, u]].unsqueeze(1) + return X + + +def digital_detect(Y, H_hat, U, N): + const = torch.tensor([1 + 1j, 1 - 1j, -1 + 1j, -1 - 1j], + device=Y.device) / np.sqrt(2.0) + 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) + Z = (torch.conj(H_hat[:, ks]) * Y[:, ks]).sum(dim=1) + metric = (Z.unsqueeze(1) * torch.conj(const).unsqueeze(0)).real + out[:, u] = metric.argmax(dim=1) + return out + + +# ------------------------------------------------------------ +# Training / adaptation / evaluation +# ------------------------------------------------------------ +DECODER_PREFIX = ("score.", "cls.") + + +def train_semantic(args, device, model_cls=BertSemanticMA, ckpt="bert_sem.pt"): + (ftr, ltr), _ = load_features(args) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + base.set_seed(args.seed) + 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 = float(rng.choice(args.train_snrs)) + fd = float(rng.choice(args.train_fds)) + x, y = sample_frames(ftr, ltr, args.batch, args.users, device, rng) + noise_var = 10 ** (-snr / 10.0) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + logits = m13.forward_frames(model, x, chan, g_p, g_d, X_pilot, + noise_var, fd, args.delta) + loss = F.cross_entropy(logits.reshape(-1, 4), y.reshape(-1)) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + if step % 1000 == 0: + print(f"[{ckpt} {step}/{args.steps}] loss={loss.item():.4f}", + flush=True) + torch.save(model.state_dict(), os.path.join(args.save_dir, ckpt)) + print("saved", ckpt) + + +def train_tx_cls(args, device): + (ftr, ltr), (fte, lte) = load_features(args) + base.set_seed(args.seed) + model = BertTxClassifier().to(device) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + for ep in range(5): + perm = torch.randperm(ftr.shape[0]) + for i in range(0, ftr.shape[0], 256): + idx = perm[i:i + 256] + logits = model(ftr[idx].to(device)) + loss = F.cross_entropy(logits, ltr[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, fte.shape[0], 2048): + acc += (model(fte[i:i + 2048].to(device)).argmax(-1) + == lte[i:i + 2048].to(device)).sum().item() + print(f"[tx-cls ep{ep + 1}] test acc={acc / fte.shape[0]:.4f}", + flush=True) + torch.save(model.state_dict(), os.path.join(args.save_dir, + "bert_txcls.pt")) + print("saved bert_txcls.pt") + + +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, x, chan, g_p, g_d, X_pilot, + noise_var, fd, delta): + y_emb, m = model.tx(x) + X_data = y_emb.to(torch.complex64) + Y_p = chan.transmit(X_pilot.unsqueeze(0).expand(x.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_bert(model, ftr, ltr, chan, X_pilot, args, snr, fd, 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(DECODER_PREFIX)} + for _ in range(args.eval_inner_steps): + x, y = sample_frames(ftr, ltr, args.support, args.users, device, rng) + g_p, g_d = chan.sample(args.support, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + with torch.enable_grad(): + logits = forward_frames_fast(model, fast, x, chan, g_p, g_d, + X_pilot, noise_var, fd, args.delta) + loss = F.cross_entropy(logits.reshape(-1, 4), y.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(args, device): + """Warm-started first-order decoder-side meta-training.""" + (ftr, ltr), _ = load_features(args) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + base.set_seed(args.seed) + model = BertSemanticMA(args.users, args.dim, args.hidden).to(device) + model.load_state_dict(torch.load(os.path.join(args.save_dir, + "bert_sem.pt"), + map_location=device)) + params = [v for k, v in model.named_parameters() + if k.startswith(DECODER_PREFIX)] + opt = torch.optim.Adam(params, lr=1e-4) + rng = np.random.default_rng(args.seed + 11) + for step in range(1, args.meta_steps + 1): + opt.zero_grad(set_to_none=True) + for _ in range(args.meta_batch): + snr = float(rng.choice(args.train_snrs)) + fd = float(rng.choice(args.train_fds)) + fast = adapt_bert(model, ftr, ltr, chan, X_pilot, args, snr, + fd, device) + fast = {k: v.requires_grad_(True) for k, v in fast.items()} + x, y = sample_frames(ftr, ltr, args.support, args.users, + device, rng) + g_p, g_d = chan.sample(args.support, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = forward_frames_fast(model, fast, x, chan, g_p, g_d, + X_pilot, noise_var, fd, args.delta) + loss = F.cross_entropy(logits.reshape(-1, 4), + y.reshape(-1)) / args.meta_batch + grads = torch.autograd.grad(loss, list(fast.values())) + named = dict(model.named_parameters()) + for (k, _), g in zip(fast.items(), grads): + if named[k].grad is None: + named[k].grad = g.detach().clone() + else: + named[k].grad += g.detach() + opt.step() + if step % 500 == 0: + print(f"[meta {step}/{args.meta_steps}]", flush=True) + torch.save(model.state_dict(), os.path.join(args.save_dir, + "bert_maml.pt")) + print("saved bert_maml.pt") + + +@torch.no_grad() +def eval_all(args, device): + base.set_seed(args.seed + 3) + (ftr, ltr), (fte, lte) = load_features(args) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + + def load(cls, name): + mdl = cls(args.users, args.dim, args.hidden).to(device) + mdl.load_state_dict(torch.load(os.path.join(args.save_dir, name), + map_location=device)) + mdl.eval() + return mdl + + sem = load(BertSemanticMA, "bert_sem.pt") + tfm = load(BertTransformerMA, "bert_tf.pt") + aem = load(BertPerUserAE, "bert_ae.pt") + mm = load(BertSemanticMA, "bert_maml.pt") + txc = BertTxClassifier().to(device) + txc.load_state_dict(torch.load(os.path.join(args.save_dir, + "bert_txcls.pt"), + map_location=device)) + txc.eval() + + rng = np.random.default_rng(args.seed + 3) + csv_path = os.path.join(args.save_dir, "bert_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, + "semantic_tf": 0, "semantic_ae": 0, "semantic_maml": 0} + fast_mm = adapt_bert(mm, ftr, ltr, chan, X_pilot, args, snr, + args.eval_fd, device) + total = 0 + for _ in range(args.eval_nb): + x, y = sample_frames(fte, lte, args.eval_batch, args.users, + device, rng) + g_p, g_d = chan.sample(args.eval_batch, args.eval_fd, + args.delta) + for key, mdl in [("semantic", sem), ("semantic_tf", tfm), + ("semantic_ae", aem)]: + lg = m13.forward_frames(mdl, x, chan, g_p, g_d, X_pilot, + noise_var, args.eval_fd, + args.delta) + errs[key] += (lg.argmax(-1) != y).sum().item() + lg = forward_frames_fast(mm, fast_mm, x, chan, g_p, g_d, + X_pilot, noise_var, args.eval_fd, + args.delta) + errs["semantic_maml"] += (lg.argmax(-1) != y).sum().item() + + B = x.shape[0] + pred = txc(x.reshape(B * args.users, -1)).argmax(-1) + pred = pred.view(B, args.users) + X_d = digital_tx(pred, 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 != y).sum().item() + total += y.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("saved", csv_path) + + +def make_fig(args): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + rows = list(csv.DictReader(open(os.path.join(args.save_dir, + "bert_results.csv")))) + 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]) + pts = [(a, b) for a, b in pts if b > 0] + if not pts: + continue + xs, ys = zip(*pts) + 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"bert_ser_vs_snr_fd{args.eval_fd}.pdf") + fig.savefig(out) + print("saved", out) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["cache", "train", "train-tf", + "train-ae", "train-tx-cls", + "train-maml", "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=10000) + 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=1500) + p.add_argument("--meta-batch", type=int, default=4) + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.01) + p.add_argument("--support", type=int, default=32) + 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("--cache-train", type=int, default=20000) + p.add_argument("--cache-test", type=int, default=7600) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--save-dir", type=str, default="results_bert") + p.add_argument("--fig-dir", type=str, default="fig") + args = p.parse_args() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("device:", device) + + if args.mode == "cache": + cache_features(args, device) + elif args.mode == "train": + train_semantic(args, device) + elif args.mode == "train-tf": + train_semantic(args, device, model_cls=BertTransformerMA, + ckpt="bert_tf.pt") + elif args.mode == "train-ae": + train_semantic(args, device, model_cls=BertPerUserAE, + ckpt="bert_ae.pt") + elif args.mode == "train-tx-cls": + train_tx_cls(args, device) + elif args.mode == "train-maml": + train_maml(args, device) + elif args.mode == "eval": + eval_all(args, device) + elif args.mode == "fig": + make_fig(args) + + +if __name__ == "__main__": + main() diff --git a/c21_mnist_flat.py b/c21_mnist_flat.py new file mode 100755 index 0000000..b4a8311 --- /dev/null +++ b/c21_mnist_flat.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +# ============================================================ +# c21_mnist_flat.py +# +# MNIST transmission over flat Rayleigh fading with AWGN, +# replacing the synthetic-symbol flat study. Compared schemes: +# Transformer SE separation, per-user AE, proposed softmax +# joint, proposed signed joint, and proposed signed MAML with +# decoder-side task-conditional adaptation (warm start from +# the signed joint model). The digital chain is +# classifier-limited on this channel and is reported from the +# transmit-side classifier accuracy. +# ============================================================ + +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 + +import c11_doppler_csi as base +import c13_mnist as m13 + + +def apply_channel(y, snr_db): + """Flat real Rayleigh + AWGN with global power normalization.""" + h = torch.randn(y.size(0), 1, device=y.device) / math.sqrt(2.0) + y = h * y + y = y / torch.sqrt(torch.mean(y ** 2) + 1e-12) + noise_var = 10 ** (-snr_db / 10.0) + return y + torch.randn_like(y) * math.sqrt(noise_var) + + +class FlatBase(nn.Module): + """CNN encoder + masks; rx defined by subclasses.""" + + def __init__(self, U, d=128, hidden=256, n_cls=10): + super().__init__() + self.U, self.d = U, d + self.encoder = m13.CNNTrunk(d) + self.masks = nn.Parameter(torch.randn(U, d)) + self.cls = nn.Sequential( + nn.Linear(d, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, n_cls)) + + def tx(self, imgs): + 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) + return y, m + + def forward(self, imgs, snr_db): + y, m = self.tx(imgs) + y = apply_channel(y, snr_db) + R = y.unsqueeze(1) * m.unsqueeze(0) + return self.rx(R) + + +class FlatSigned(FlatBase): + def __init__(self, U, d=128, hidden=256, n_cls=10, score_hidden=64): + super().__init__(U, d, hidden, n_cls) + 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)) + + def rx(self, R): + B = R.shape[0] + T = torch.bmm(R, R.transpose(1, 2)) / self.d + w = self.score(T.reshape(B, -1)).view(B, self.U, self.U) + W = torch.eye(self.U, device=R.device).unsqueeze(0) + w + return self.cls(torch.bmm(W, R)) + + +class FlatSoftmax(FlatBase): + def __init__(self, U, d=128, hidden=256, n_cls=10): + super().__init__(U, d, hidden, n_cls) + self.query = nn.Parameter(torch.randn(U, d)) + self.key = nn.Linear(d, d, bias=False) + self.val = nn.Linear(d, d, bias=False) + + def rx(self, R): + K, Vv = self.key(R), self.val(R) + q = F.normalize(self.query, dim=1) + scores = torch.einsum("ud,bid->bui", q, + F.normalize(K, dim=-1)) / math.sqrt(self.d) + attn = F.softmax(scores * 1.43, dim=-1) + z = torch.einsum("bui,bid->bud", attn, Vv) + R + return self.cls(z) + + +class FlatTransformer(FlatBase): + def __init__(self, U, d=128, hidden=256, n_cls=10, n_heads=4, + n_layers=2): + super().__init__(U, d, hidden, n_cls) + layer = nn.TransformerEncoderLayer(d_model=hidden, nhead=n_heads, + dim_feedforward=2 * hidden, + batch_first=True) + self.inp = nn.Linear(d, hidden) + self.sep = nn.TransformerEncoder(layer, num_layers=n_layers) + self.out = nn.Linear(hidden, n_cls) + + def rx(self, R): + return self.out(self.sep(self.inp(R))) + + +class FlatPerUserAE(nn.Module): + """Per-user CNN encoders and heads, no masks.""" + + def __init__(self, U, d=128, hidden=256, n_cls=10): + super().__init__() + self.U, self.d = U, d + self.encs = nn.ModuleList([m13.CNNTrunk(d) for _ in range(U)]) + self.heads = nn.ModuleList([ + nn.Sequential(nn.Linear(d, hidden), nn.ReLU(inplace=True), + nn.Linear(hidden, n_cls)) + for _ in range(U) + ]) + + def forward(self, imgs, snr_db): + e = torch.stack([self.encs[u](imgs[:, u]) for u in range(self.U)], + dim=1) + y = e.sum(dim=1) + y = apply_channel(y, snr_db) + return torch.stack([h(y) for h in self.heads], dim=1) + + +MODELS = { + "transformer": FlatTransformer, + "peruser": FlatPerUserAE, + "softmax": FlatSoftmax, + "signed": FlatSigned, +} + +DECODER_PREFIX = ("score.", "cls.") + + +def train_model(key, args, device, x_tr, y_tr): + base.set_seed(args.seed) + model = MODELS[key](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 = float(rng.choice(args.train_snrs)) + imgs, labels = m13.sample_frames(x_tr, y_tr, args.batch, args.users, + device, rng) + logits = model(imgs, snr) + loss = F.cross_entropy(logits.reshape(-1, 10), labels.reshape(-1)) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + if step % 2000 == 0: + print(f"[{key} {step}/{args.steps}] loss={loss.item():.4f}", + flush=True) + return model + + +def rx_with_fast(model, R, fast): + B = R.shape[0] + T = torch.bmm(R, R.transpose(1, 2)) / model.d + h1 = torch.relu(F.linear(T.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"])) + w = F.linear(h1, fast["score.4.weight"], fast["score.4.bias"]) + W = torch.eye(model.U, device=R.device).unsqueeze(0) + W = W + w.view(B, model.U, model.U) + z = torch.bmm(W, R) + 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 adapt_decoder(model, x_tr, y_tr, 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(DECODER_PREFIX)} + for _ in range(args.eval_inner_steps): + imgs, labels = m13.sample_frames(x_tr, y_tr, args.support, + args.users, device, rng) + with torch.enable_grad(): + y, m = model.tx(imgs) + y = apply_channel(y, snr) + R = y.unsqueeze(1) * m.unsqueeze(0) + logits = rx_with_fast(model, R, fast) + 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 meta_train_decoder(model, x_tr, y_tr, args, device): + """First-order MAML on the decoder-side parameters, warm-started + from the jointly trained signed model.""" + params = [v for k, v in model.named_parameters() + if k.startswith(DECODER_PREFIX)] + opt = torch.optim.Adam(params, lr=1e-4) + rng = np.random.default_rng(args.seed + 11) + for step in range(1, args.meta_steps + 1): + opt.zero_grad(set_to_none=True) + for _ in range(args.meta_batch): + snr = float(rng.choice(args.train_snrs)) + fast = adapt_decoder(model, x_tr, y_tr, args, snr, device) + fast = {k: v.requires_grad_(True) for k, v in fast.items()} + imgs, labels = m13.sample_frames(x_tr, y_tr, args.support, + args.users, device, rng) + ytx, m = model.tx(imgs) + ytx = apply_channel(ytx, snr) + R = ytx.unsqueeze(1) * m.unsqueeze(0) + logits = rx_with_fast(model, R, fast) + loss = F.cross_entropy(logits.reshape(-1, 10), + labels.reshape(-1)) / args.meta_batch + grads = torch.autograd.grad(loss, list(fast.values())) + named = dict(model.named_parameters()) + for (k, _), g in zip(fast.items(), grads): + if named[k].grad is None: + named[k].grad = g.detach().clone() + else: + named[k].grad += g.detach() + opt.step() + if step % 500 == 0: + print(f"[meta {step}/{args.meta_steps}]", flush=True) + return model + + +@torch.no_grad() +def eval_model(model, x, y, args, snr, device, rng, fast=None): + errs, total = 0, 0 + for _ in range(args.eval_nb): + imgs, labels = m13.sample_frames(x, y, args.eval_batch, args.users, + device, rng) + if fast is None: + logits = model(imgs, snr) + else: + ytx, m = model.tx(imgs) + ytx = apply_channel(ytx, snr) + R = ytx.unsqueeze(1) * m.unsqueeze(0) + logits = rx_with_fast(model, R, fast) + errs += (logits.argmax(-1) != labels).sum().item() + total += labels.numel() + return errs / total + + +def eval_maml_only(args, device): + """Precision re-evaluation of the flat signed MAML receiver.""" + tr, te = m13.get_datasets(args.data_root) + x_tr, y_tr = m13.tensorize(tr) + x_te, y_te = m13.tensorize(te) + model = FlatSigned(args.users, args.dim, args.hidden).to(device) + model.load_state_dict(torch.load( + os.path.join(args.save_dir, "flat_maml.pt"), map_location=device)) + model.eval() + joint = FlatSigned(args.users, args.dim, args.hidden).to(device) + joint.load_state_dict(torch.load( + os.path.join(args.save_dir, "flat_signed.pt"), map_location=device)) + joint.eval() + for snr in args.eval_snrs: + sers = [] + for rep in range(args.adapt_reps): + args_seed = args.seed + 77 + rep + rng = np.random.default_rng(args_seed) + fast = {k: v.detach().clone().requires_grad_(True) + for k, v in model.named_parameters() + if k.startswith(DECODER_PREFIX)} + for _ in range(args.eval_inner_steps): + imgs, labels = m13.sample_frames(x_tr, y_tr, args.support, + args.users, device, rng) + with torch.enable_grad(): + y, m = model.tx(imgs) + y = apply_channel(y, float(snr)) + R = y.unsqueeze(1) * m.unsqueeze(0) + logits = rx_with_fast(model, R, fast) + 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)} + fast = {k: v.detach() for k, v in fast.items()} + ser = eval_model(model, x_te, y_te, args, float(snr), device, + np.random.default_rng(args.seed + 3), + fast=fast) + sers.append(ser) + sj = eval_model(joint, x_te, y_te, args, float(snr), device, + np.random.default_rng(args.seed + 3)) + print(f"snr={snr}: joint={sj:.4e} maml mean={np.mean(sers):.4e} " + f"reps={[f'{s:.4e}' for s in sers]}", flush=True) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["run", "eval-maml"], default="run") + p.add_argument("--adapt-reps", type=int, default=5) + 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("--train-snrs", type=float, nargs="+", + default=[0, 5, 10, 15, 20, 25, 30]) + p.add_argument("--steps", type=int, default=10000) + p.add_argument("--batch", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--eval-snrs", type=float, nargs="+", + default=[10, 20, 30]) + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.01) + p.add_argument("--meta-steps", type=int, default=1500) + p.add_argument("--meta-batch", type=int, default=4) + p.add_argument("--support", type=int, default=32) + 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("--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 == "eval-maml": + eval_maml_only(args, device) + return + + tr, te = m13.get_datasets(args.data_root) + x_tr, y_tr = m13.tensorize(tr) + x_te, y_te = m13.tensorize(te) + + csv_path = os.path.join(args.save_dir, "mnist_flat.csv") + with open(csv_path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["method", "snr_db", "ser"]) + signed_model = None + for key in MODELS: + model = train_model(key, args, device, x_tr, y_tr) + if key == "signed": + signed_model = model + torch.save(model.state_dict(), + os.path.join(args.save_dir, "flat_signed.pt")) + for snr in args.eval_snrs: + ser = eval_model(model, x_te, y_te, args, float(snr), + device, np.random.default_rng(args.seed + 3)) + w.writerow([key, snr, ser]) + f.flush() + print(f"[{key}] snr={snr} ser={ser:.4e}", flush=True) + # signed MAML: warm-started decoder-side meta-training, then + # task-conditional adaptation per evaluation SNR + signed_model = meta_train_decoder(signed_model, x_tr, y_tr, args, + device) + torch.save(signed_model.state_dict(), + os.path.join(args.save_dir, "flat_maml.pt")) + for snr in args.eval_snrs: + fast = adapt_decoder(signed_model, x_tr, y_tr, args, + float(snr), device) + ser = eval_model(signed_model, x_te, y_te, args, float(snr), + device, np.random.default_rng(args.seed + 3), + fast=fast) + w.writerow(["signed_maml", snr, ser]) + f.flush() + print(f"[signed_maml] snr={snr} ser={ser:.4e}", flush=True) + print("saved", csv_path) + + +if __name__ == "__main__": + main() diff --git a/c22_maml_epoch.py b/c22_maml_epoch.py new file mode 100755 index 0000000..ec24e50 --- /dev/null +++ b/c22_maml_epoch.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +# ============================================================ +# c22_maml_epoch.py +# +# Meta-training trajectory for the signed MAML receiver on the +# MNIST TDL chain: warm-started from the converged joint model, +# the adapted SER at (15 dB, nu = 0.05) is recorded every +# ckpt-every meta-steps, extending the convergence figure of +# c19 beyond the joint-training phase. +# ============================================================ + +import argparse +import csv +import os +import numpy as np +import torch +import torch.nn.functional as F + +import c11_doppler_csi as base +import c13_mnist as m13 + + +@torch.no_grad() +def eval_adapted(model, fast, x, y, chan, X_pilot, args, device): + noise_var = 10 ** (-args.eval_snr / 10.0) + rng = np.random.default_rng(args.seed + 3) + errs, total = 0, 0 + for _ in range(args.eval_nb): + imgs, labels = m13.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) + logits = m13.forward_frames_fast(model, fast, imgs, chan, g_p, g_d, + X_pilot, noise_var, args.eval_fd, + args.delta) + errs += (logits.argmax(-1) != labels).sum().item() + total += labels.numel() + return errs / total + + +def main(): + p = argparse.ArgumentParser() + 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("--meta-steps", type=int, default=2500) + p.add_argument("--pretrain-steps", type=int, default=0) + p.add_argument("--ckpt-every", type=int, default=250) + p.add_argument("--meta-batch", type=int, default=4) + p.add_argument("--inner-lr", type=float, default=0.02) + p.add_argument("--batch", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--eval-inner-steps", type=int, default=5) + p.add_argument("--eval-inner-lr", type=float, default=0.01) + p.add_argument("--support", type=int, default=32) + p.add_argument("--eval-snr", type=float, default=15.0) + 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=30) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--save-dir", type=str, default="results_mnist") + 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) + + base.set_seed(args.seed) + tr, te = m13.get_datasets(args.data_root) + x_tr, y_tr = m13.tensorize(tr) + x_te, y_te = m13.tensorize(te) + chan = base.TDLChannel(args.nfft, args.cp, args.taps, device=device) + X_pilot = base.make_pilot(args.nfft, device) + model = m13.MnistSemanticMA(args.users, args.dim, args.hidden).to(device) + rng = np.random.default_rng(args.seed) + if args.pretrain_steps > 0: + # retrace the joint training for the warm-start phase so that + # the meta phase continues the same trajectory as the joint curve + popt = torch.optim.Adam(model.parameters(), lr=1e-3) + for step in range(1, args.pretrain_steps + 1): + snr = float(rng.choice(args.train_snrs)) + fd = float(rng.choice(args.train_fds)) + imgs, labels = m13.sample_frames(x_tr, y_tr, args.batch, + args.users, device, rng) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = m13.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)) + popt.zero_grad(set_to_none=True) + loss.backward() + popt.step() + if step % 2000 == 0: + print(f"[pretrain {step}/{args.pretrain_steps}]", flush=True) + else: + model.load_state_dict(torch.load( + os.path.join(args.save_dir, "mnist_semantic.pt"), + map_location=device)) + opt = torch.optim.Adam(model.parameters(), lr=args.lr) + + def task_loss(fast, snr, fd): + imgs, labels = m13.sample_frames(x_tr, y_tr, args.batch, args.users, + device, rng) + g_p, g_d = chan.sample(args.batch, fd, args.delta) + noise_var = 10 ** (-snr / 10.0) + logits = m13.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)) + + csv_path = os.path.join(args.save_dir, "mnist_maml_epoch.csv") + with open(csv_path, "w", newline="") as f: + w = csv.writer(f) + w.writerow(["step", "ser"]) + fast0 = m13.adapt_mnist(model, x_tr, y_tr, chan, X_pilot, args, + args.eval_snr, device) + ser0 = eval_adapted(model, fast0, x_te, y_te, chan, X_pilot, args, + device) + w.writerow([args.pretrain_steps, ser0]) + print(f"[meta 0] ser={ser0:.4e}", flush=True) + 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(m13.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 % args.ckpt_every == 0: + fastc = m13.adapt_mnist(model, x_tr, y_tr, chan, X_pilot, + args, args.eval_snr, device) + ser = eval_adapted(model, fastc, x_te, y_te, chan, X_pilot, + args, device) + w.writerow([args.pretrain_steps + step, ser]) + f.flush() + print(f"[meta {step}/{args.meta_steps}] ser={ser:.4e}", + flush=True) + print("saved", csv_path) + + +if __name__ == "__main__": + main() diff --git a/fig/bert_ser_vs_snr_fd0.05.pdf b/fig/bert_ser_vs_snr_fd0.05.pdf new file mode 100755 index 0000000..827b096 Binary files /dev/null and b/fig/bert_ser_vs_snr_fd0.05.pdf differ diff --git a/fig/mnist_epoch_convergence.pdf b/fig/mnist_epoch_convergence.pdf new file mode 100755 index 0000000..7c5ea49 Binary files /dev/null and b/fig/mnist_epoch_convergence.pdf differ diff --git a/fig/mnist_ser_vs_doppler_snr15.pdf b/fig/mnist_ser_vs_doppler_snr15.pdf new file mode 100755 index 0000000..15fb2d0 Binary files /dev/null and b/fig/mnist_ser_vs_doppler_snr15.pdf differ diff --git a/fig/mnist_ser_vs_snr_fd0.05.pdf b/fig/mnist_ser_vs_snr_fd0.05.pdf new file mode 100755 index 0000000..132780f Binary files /dev/null and b/fig/mnist_ser_vs_snr_fd0.05.pdf differ diff --git a/results_bert/bert_ae.pt b/results_bert/bert_ae.pt new file mode 100755 index 0000000..20b8846 Binary files /dev/null and b/results_bert/bert_ae.pt differ diff --git a/results_bert/bert_maml.pt b/results_bert/bert_maml.pt new file mode 100755 index 0000000..ef866fe Binary files /dev/null and b/results_bert/bert_maml.pt differ diff --git a/results_bert/bert_results.csv b/results_bert/bert_results.csv new file mode 100755 index 0000000..11eecc2 --- /dev/null +++ b/results_bert/bert_results.csv @@ -0,0 +1,43 @@ +snr_db,fd_norm,method,ser +0,0.05,digital_genie,0.11959635416666667 +0,0.05,digital_pilot,0.6495442708333333 +0,0.05,semantic,0.38662109375 +0,0.05,semantic_tf,0.42034505208333334 +0,0.05,semantic_ae,0.40579427083333336 +0,0.05,semantic_maml,0.37786458333333334 +5,0.05,digital_genie,0.115625 +5,0.05,digital_pilot,0.6171549479166667 +5,0.05,semantic,0.20755208333333333 +5,0.05,semantic_tf,0.248046875 +5,0.05,semantic_ae,0.2296875 +5,0.05,semantic_maml,0.20283203125 +10,0.05,digital_genie,0.1119140625 +10,0.05,digital_pilot,0.603125 +10,0.05,semantic,0.15341796875 +10,0.05,semantic_tf,0.18011067708333334 +10,0.05,semantic_ae,0.16171875 +10,0.05,semantic_maml,0.1482421875 +15,0.05,digital_genie,0.11533203125 +15,0.05,digital_pilot,0.6046875 +15,0.05,semantic,0.13987630208333332 +15,0.05,semantic_tf,0.16845703125 +15,0.05,semantic_ae,0.15110677083333332 +15,0.05,semantic_maml,0.13935546875 +20,0.05,digital_genie,0.11744791666666667 +20,0.05,digital_pilot,0.59267578125 +20,0.05,semantic,0.13942057291666668 +20,0.05,semantic_tf,0.16389973958333334 +20,0.05,semantic_ae,0.14912109375 +20,0.05,semantic_maml,0.13756510416666667 +25,0.05,digital_genie,0.11673177083333333 +25,0.05,digital_pilot,0.60029296875 +25,0.05,semantic,0.1400390625 +25,0.05,semantic_tf,0.16061197916666667 +25,0.05,semantic_ae,0.14612630208333333 +25,0.05,semantic_maml,0.13909505208333334 +30,0.05,digital_genie,0.1162109375 +30,0.05,digital_pilot,0.6109049479166667 +30,0.05,semantic,0.14016927083333333 +30,0.05,semantic_tf,0.16246744791666667 +30,0.05,semantic_ae,0.14560546875 +30,0.05,semantic_maml,0.13753255208333334 diff --git a/results_bert/bert_sem.pt b/results_bert/bert_sem.pt new file mode 100755 index 0000000..c8dd638 Binary files /dev/null and b/results_bert/bert_sem.pt differ diff --git a/results_bert/bert_tf.pt b/results_bert/bert_tf.pt new file mode 100755 index 0000000..5f4917a Binary files /dev/null and b/results_bert/bert_tf.pt differ diff --git a/results_bert/bert_txcls.pt b/results_bert/bert_txcls.pt new file mode 100755 index 0000000..4953e7a Binary files /dev/null and b/results_bert/bert_txcls.pt differ diff --git a/results_mnist/flat_maml.pt b/results_mnist/flat_maml.pt new file mode 100755 index 0000000..3033790 Binary files /dev/null and b/results_mnist/flat_maml.pt differ diff --git a/results_mnist/flat_signed.pt b/results_mnist/flat_signed.pt new file mode 100755 index 0000000..a3b4bee Binary files /dev/null and b/results_mnist/flat_signed.pt differ diff --git a/results_mnist/mnist_ae.pt b/results_mnist/mnist_ae.pt new file mode 100755 index 0000000..c6be5ce Binary files /dev/null and b/results_mnist/mnist_ae.pt differ diff --git a/results_mnist/mnist_doppler.csv b/results_mnist/mnist_doppler.csv new file mode 100755 index 0000000..883d570 --- /dev/null +++ b/results_mnist/mnist_doppler.csv @@ -0,0 +1,67 @@ +snr_db,fd_norm,method,ser +15.0,0.001,digital_genie,0.021614583333333333 +15.0,0.001,digital_pilot,0.021647135416666668 +15.0,0.001,semantic,0.017122395833333335 +15.0,0.001,semantic_tf,0.01943359375 +15.0,0.001,semantic_ae,0.027376302083333335 +15.0,0.001,semantic_maml,0.015983072916666667 +15.0,0.002,digital_genie,0.020377604166666667 +15.0,0.002,digital_pilot,0.020377604166666667 +15.0,0.002,semantic,0.016634114583333335 +15.0,0.002,semantic_tf,0.017220052083333333 +15.0,0.002,semantic_ae,0.024544270833333333 +15.0,0.002,semantic_maml,0.013834635416666666 +15.0,0.005,digital_genie,0.02138671875 +15.0,0.005,digital_pilot,0.02138671875 +15.0,0.005,semantic,0.017350260416666666 +15.0,0.005,semantic_tf,0.019368489583333332 +15.0,0.005,semantic_ae,0.02705078125 +15.0,0.005,semantic_maml,0.01552734375 +15.0,0.01,digital_genie,0.02255859375 +15.0,0.01,digital_pilot,0.02373046875 +15.0,0.01,semantic,0.02158203125 +15.0,0.01,semantic_tf,0.022623697916666668 +15.0,0.01,semantic_ae,0.030078125 +15.0,0.01,semantic_maml,0.01845703125 +15.0,0.02,digital_genie,0.022623697916666668 +15.0,0.02,digital_pilot,0.07255859375 +15.0,0.02,semantic,0.026627604166666666 +15.0,0.02,semantic_tf,0.03382161458333333 +15.0,0.02,semantic_ae,0.042740885416666666 +15.0,0.02,semantic_maml,0.02236328125 +15.0,0.03,digital_genie,0.0208984375 +15.0,0.03,digital_pilot,0.28782552083333335 +15.0,0.03,semantic,0.04156901041666667 +15.0,0.03,semantic_tf,0.0587890625 +15.0,0.03,semantic_ae,0.06376953125 +15.0,0.03,semantic_maml,0.031510416666666666 +15.0,0.05,digital_genie,0.021451822916666665 +15.0,0.05,digital_pilot,0.8395182291666666 +15.0,0.05,semantic,0.07496744791666667 +15.0,0.05,semantic_tf,0.11764322916666667 +15.0,0.05,semantic_ae,0.13313802083333334 +15.0,0.05,semantic_maml,0.04938151041666667 +15.0,0.0567,digital_genie,0.022005208333333335 +15.0,0.0567,digital_pilot,0.9119466145833334 +15.0,0.0567,semantic,0.0775390625 +15.0,0.0567,semantic_tf,0.12428385416666667 +15.0,0.0567,semantic_ae,0.13880208333333333 +15.0,0.0567,semantic_maml,0.051953125 +15.0,0.07,digital_genie,0.022884114583333334 +15.0,0.07,digital_pilot,0.9771484375 +15.0,0.07,semantic,0.07864583333333333 +15.0,0.07,semantic_tf,0.115625 +15.0,0.07,semantic_ae,0.13460286458333334 +15.0,0.07,semantic_maml,0.05462239583333333 +15.0,0.085,digital_genie,0.02041015625 +15.0,0.085,digital_pilot,0.9912760416666667 +15.0,0.085,semantic,0.06474609375 +15.0,0.085,semantic_tf,0.09527994791666666 +15.0,0.085,semantic_ae,0.11110026041666667 +15.0,0.085,semantic_maml,0.045572916666666664 +15.0,0.1,digital_genie,0.020833333333333332 +15.0,0.1,digital_pilot,0.9902018229166667 +15.0,0.1,semantic,0.0673828125 +15.0,0.1,semantic_tf,0.10113932291666666 +15.0,0.1,semantic_ae,0.1130859375 +15.0,0.1,semantic_maml,0.04791666666666667 diff --git a/results_mnist/mnist_epoch.csv b/results_mnist/mnist_epoch.csv new file mode 100755 index 0000000..c9e3648 --- /dev/null +++ b/results_mnist/mnist_epoch.csv @@ -0,0 +1,80 @@ +method,step,ser +semantic,0,0.9009765625 +semantic,500,0.6244140625 +semantic,1000,0.5520182291666667 +semantic,1500,0.4524739583333333 +semantic,2000,0.33196614583333334 +semantic,2500,0.259765625 +semantic,3000,0.21471354166666667 +semantic,3500,0.19192708333333333 +semantic,4000,0.17024739583333334 +semantic,4500,0.14934895833333334 +semantic,5000,0.1462890625 +semantic,5500,0.13203125 +semantic,6000,0.11979166666666667 +semantic,6500,0.11809895833333334 +semantic,7000,0.10872395833333333 +semantic,7500,0.10078125 +semantic,8000,0.0984375 +semantic,8500,0.09563802083333334 +semantic,9000,0.0904296875 +semantic,9500,0.084765625 +semantic,10000,0.07994791666666666 +semantic_tf,0,0.8890625 +semantic_tf,500,0.63671875 +semantic_tf,1000,0.56953125 +semantic_tf,1500,0.47630208333333335 +semantic_tf,2000,0.38326822916666664 +semantic_tf,2500,0.33587239583333334 +semantic_tf,3000,0.3083984375 +semantic_tf,3500,0.2765625 +semantic_tf,4000,0.245703125 +semantic_tf,4500,0.21204427083333333 +semantic_tf,5000,0.18294270833333334 +semantic_tf,5500,0.17838541666666666 +semantic_tf,6000,0.15852864583333334 +semantic_tf,6500,0.158984375 +semantic_tf,7000,0.1466796875 +semantic_tf,7500,0.14733072916666667 +semantic_tf,8000,0.13990885416666668 +semantic_tf,8500,0.12350260416666667 +semantic_tf,9000,0.1166015625 +semantic_tf,9500,0.12024739583333334 +semantic_tf,10000,0.10501302083333333 +semantic_ae,0,0.9024739583333333 +semantic_ae,500,0.658203125 +semantic_ae,1000,0.6228515625 +semantic_ae,1500,0.6057291666666667 +semantic_ae,2000,0.5608723958333334 +semantic_ae,2500,0.503515625 +semantic_ae,3000,0.4469401041666667 +semantic_ae,3500,0.4055989583333333 +semantic_ae,4000,0.34205729166666665 +semantic_ae,4500,0.2892578125 +semantic_ae,5000,0.2609375 +semantic_ae,5500,0.22161458333333334 +semantic_ae,6000,0.20826822916666668 +semantic_ae,6500,0.19466145833333334 +semantic_ae,7000,0.17571614583333334 +semantic_ae,7500,0.15787760416666666 +semantic_ae,8000,0.15143229166666666 +semantic_ae,8500,0.1470703125 +semantic_ae,9000,0.13880208333333333 +semantic_ae,9500,0.13404947916666668 +semantic_ae,10000,0.12194010416666666 +semantic_maml_pre,0,0.901171875 +semantic_maml_pre,500,0.6270182291666667 +semantic_maml_pre,1000,0.558984375 +semantic_maml_pre,1500,0.4520182291666667 +semantic_maml_pre,2000,0.3409505208333333 +semantic_maml_pre,2500,0.27239583333333334 +semantic_maml_pre,3000,0.22682291666666668 +semantic_maml_pre,3500,0.21061197916666666 +semantic_maml_pre,4000,0.18743489583333334 +semantic_maml_pre,4500,0.1646484375 +semantic_maml_pre,5000,0.14342447916666667 +semantic_maml_pre,5500,0.13248697916666666 +semantic_maml_pre,6000,0.122265625 +semantic_maml_pre,6500,0.11985677083333333 +semantic_maml_pre,7000,0.10904947916666667 +semantic_maml_pre,7500,0.10865885416666667 diff --git a/results_mnist/mnist_flat.csv b/results_mnist/mnist_flat.csv new file mode 100755 index 0000000..00a88d3 --- /dev/null +++ b/results_mnist/mnist_flat.csv @@ -0,0 +1,16 @@ +method,snr_db,ser +transformer,10,0.14541015625 +transformer,20,0.05712890625 +transformer,30,0.029850260416666666 +peruser,10,0.16139322916666668 +peruser,20,0.06409505208333334 +peruser,30,0.037760416666666664 +softmax,10,0.15491536458333333 +softmax,20,0.058984375 +softmax,30,0.03313802083333333 +signed,10,0.14690755208333334 +signed,20,0.05218098958333333 +signed,30,0.024348958333333334 +signed_maml,10,0.13597005208333332 +signed_maml,20,0.0509765625 +signed_maml,30,0.026822916666666665 diff --git a/results_mnist/mnist_maml.pt b/results_mnist/mnist_maml.pt new file mode 100755 index 0000000..801ad62 Binary files /dev/null and b/results_mnist/mnist_maml.pt differ diff --git a/results_mnist/mnist_maml_epoch.csv b/results_mnist/mnist_maml_epoch.csv new file mode 100755 index 0000000..bbfc938 --- /dev/null +++ b/results_mnist/mnist_maml_epoch.csv @@ -0,0 +1,12 @@ +step,ser +7500,0.1025390625 +7750,0.09036458333333333 +8000,0.08580729166666666 +8250,0.07532552083333334 +8500,0.07350260416666667 +8750,0.07864583333333333 +9000,0.06575520833333333 +9250,0.0638671875 +9500,0.05865885416666667 +9750,0.0591796875 +10000,0.05755208333333333 diff --git a/results_mnist/mnist_results.csv b/results_mnist/mnist_results.csv new file mode 100755 index 0000000..5e95607 --- /dev/null +++ b/results_mnist/mnist_results.csv @@ -0,0 +1,43 @@ +snr_db,fd_norm,method,ser +0,0.05,digital_genie,0.06774088541666666 +0,0.05,digital_pilot,0.8790364583333333 +0,0.05,semantic,0.5509765625 +0,0.05,semantic_tf,0.5806315104166667 +0,0.05,semantic_ae,0.6121744791666667 +0,0.05,semantic_maml,0.5267578125 +5,0.05,digital_genie,0.022102864583333333 +5,0.05,digital_pilot,0.8439778645833333 +5,0.05,semantic,0.25537109375 +5,0.05,semantic_tf,0.30462239583333334 +5,0.05,semantic_ae,0.3328776041666667 +5,0.05,semantic_maml,0.21705729166666668 +10,0.05,digital_genie,0.02138671875 +10,0.05,digital_pilot,0.83134765625 +10,0.05,semantic,0.11331380208333333 +10,0.05,semantic_tf,0.1607421875 +10,0.05,semantic_ae,0.17711588541666667 +10,0.05,semantic_maml,0.07799479166666666 +15,0.05,digital_genie,0.02255859375 +15,0.05,digital_pilot,0.8248372395833333 +15,0.05,semantic,0.07705078125 +15,0.05,semantic_tf,0.11637369791666667 +15,0.05,semantic_ae,0.12939453125 +15,0.05,semantic_maml,0.0521484375 +20,0.05,digital_genie,0.022623697916666668 +20,0.05,digital_pilot,0.82109375 +20,0.05,semantic,0.064453125 +20,0.05,semantic_tf,0.10296223958333334 +20,0.05,semantic_ae,0.1150390625 +20,0.05,semantic_maml,0.04309895833333333 +25,0.05,digital_genie,0.0208984375 +25,0.05,digital_pilot,0.8190755208333333 +25,0.05,semantic,0.06419270833333333 +25,0.05,semantic_tf,0.09918619791666666 +25,0.05,semantic_ae,0.11041666666666666 +25,0.05,semantic_maml,0.04358723958333333 +30,0.05,digital_genie,0.021451822916666665 +30,0.05,digital_pilot,0.8393880208333333 +30,0.05,semantic,0.05862630208333333 +30,0.05,semantic_tf,0.09840494791666667 +30,0.05,semantic_ae,0.11266276041666666 +30,0.05,semantic_maml,0.0380859375 diff --git a/results_mnist/mnist_semantic.pt b/results_mnist/mnist_semantic.pt new file mode 100755 index 0000000..0684520 Binary files /dev/null and b/results_mnist/mnist_semantic.pt differ diff --git a/results_mnist/mnist_tf.pt b/results_mnist/mnist_tf.pt new file mode 100755 index 0000000..1946b7d Binary files /dev/null and b/results_mnist/mnist_tf.pt differ diff --git a/results_mnist/mnist_txcls.pt b/results_mnist/mnist_txcls.pt new file mode 100755 index 0000000..bc9b432 Binary files /dev/null and b/results_mnist/mnist_txcls.pt differ