156 lines
7.0 KiB
Python
Executable File
156 lines
7.0 KiB
Python
Executable File
#!/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()
|