Code and stored results for AI-native multi-user semantic communications (JSAC submission)
This commit is contained in:
KiHoLee
Executable
+369
@@ -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()
|
||||
Reference in New Issue
Block a user