Hugging Face's logo

wi-lab
/
lwm-cde

PyTorch
wireless
large-wireless-model
foundation-model
dataset-distance
transfer-learning
deepmimo
Model card Files
xet
 Community
lwm-cde / tutorial.py
wi-lab's picture
wi-lab
LWM-CDE: dataset distancing for wireless datasets
7378252 verified
Raw History Blame Contribute Delete
14.8 kB
"""
Dataset Distancing Tutorial
============================
End-to-end: generate datasets from DeepMIMO, split into train/val/test,
compute distances and transfer performance, measure correlation.
Usage:
python tutorial.py
python tutorial.py --trials 5
python tutorial.py --task beam-pred --n_beams 8 --epochs 150 --batch_size 512
"""
import argparse, logging, sys, os
parser = argparse.ArgumentParser(description="Dataset Distancing Tutorial")
# paths
parser.add_argument("--lwm_path", type=str, default="third_party/lwm-1.1")
parser.add_argument("--weights_file", type=str, default=None, help="alternative checkpoint (e.g. CDE-finetuned)")
parser.add_argument("--cache_dir", type=str, default="cache/tutorial")
parser.add_argument("--gpu", type=int, default=1)
# data splits
parser.add_argument("--n_train", type=int, default=3000)
parser.add_argument("--n_val", type=int, default=2000)
parser.add_argument("--n_test", type=int, default=2000)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trials", type=int, default=1, help="number of trials with different split seeds")
# scenario generation
parser.add_argument("--task", type=str, default="los-nlos",
choices=["los-nlos", "beam-pred"])
parser.add_argument("--n_ant", type=int, default=32)
parser.add_argument("--n_sub", type=int, default=32)
parser.add_argument("--n_beams", type=int, default=32, help="codebook size for beam-pred")
# distance computation
parser.add_argument("--distance", type=str, default="swd", choices=["swd", "centroid", "cosine"])
parser.add_argument("--embed_space", type=str, default="lwm", choices=["lwm", "raw", "umap"],
help="embedding space for distance computation")
parser.add_argument("--umap_dim", type=int, default=2, help="UMAP target dimension")
parser.add_argument("--label_aware", action="store_true", default=True)
parser.add_argument("--no_label_aware", action="store_true")
parser.add_argument("--sw_projections", type=int, default=64)
parser.add_argument("--asymmetric", action="store_true", help="directed SWD + unsymmetrized P")
parser.add_argument("--raw_perf", action="store_true", help="skip P symmetrization (keep symmetric D)")
parser.add_argument("--asym_alpha", type=float, default=1.0, help="strength of asymmetry")
# CDE fine-tuning
parser.add_argument("--cde", action="store_true", help="train CDE before computing distances")
parser.add_argument("--cde_epochs", type=int, default=40)
parser.add_argument("--cde_lr", type=float, default=5e-4)
parser.add_argument("--cde_save", type=str, default=None, help="save CDE checkpoint to this dir")
# performance matrix
parser.add_argument("--epochs", type=int, default=15)
parser.add_argument("--lr", type=float, default=1e-3)
parser.add_argument("--batch_size", type=int, default=128)
# label robustness
parser.add_argument("--label_noise", type=float, default=0.0, help="fraction of labels to flip randomly")
parser.add_argument("--label_drop", type=float, default=0.0, help="fraction of labels to drop (β†’ label-agnostic D)")
parser.add_argument("--log_dir", type=str, default="logs", help="directory for log files")
args = parser.parse_args()
if args.no_label_aware:
args.label_aware = False
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
# ── Logging setup ────────────────────────────────────────────────────────────
os.makedirs(args.log_dir, exist_ok=True)
from datetime import datetime
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = os.path.join(args.log_dir, f"{args.task}_{timestamp}.log")
log = logging.getLogger("tutorial")
log.setLevel(logging.INFO)
log.handlers.clear()
# File handler
fh = logging.FileHandler(log_file)
fh.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%H:%M:%S"))
log.addHandler(fh)
# Console handler
ch = logging.StreamHandler()
ch.setFormatter(logging.Formatter("%(message)s"))
log.addHandler(ch)
log.info(f"Command: {' '.join(sys.argv)}")
log.info(f"Args: {vars(args)}")
log.info(f"Log file: {log_file}")
import torch, numpy as np
from dataset_distancing import (
DatasetCollection,
compute_distances,
compute_performance_matrix,
correlate,
DistanceConfig,
TransferConfig,
)
from dataset_distancing.embeddings import register_lwm_path, load_runner
register_lwm_path(args.lwm_path)
# ── Step 1: Generate datasets (once, shared across trials) ───────────────────
import deepmimo as dm
scenarios = (
[{"name": s, "tx_set": tx}
for s in dm.search()[1:5]
for tx in [1]]
+ dm.search()[60:65]
)
full = DatasetCollection.from_scenarios(
scenarios,
task=args.task,
n_ant=args.n_ant,
n_sub=args.n_sub,
n_beams=args.n_beams,
cache_dir=args.cache_dir,
)
def _print_matrix(name, M):
log.info(f"\n{name}:")
for row in M:
log.info(" " + " ".join(f"{v:.4f}" for v in row))
def _perturb_labels(collection, noise_frac, drop_frac, seed):
"""Flip and/or drop a fraction of labels in each dataset."""
from dataset_distancing.data import WirelessDataset
g = torch.Generator().manual_seed(seed)
new_ds = []
for ds in collection.datasets:
if ds.labels is None:
new_ds.append(ds)
continue
labels = ds.labels.clone()
n = len(labels)
# Flip labels
if noise_frac > 0:
n_flip = int(noise_frac * n)
flip_idx = torch.randperm(n, generator=g)[:n_flip]
n_classes = int(labels.max().item()) + 1
labels[flip_idx] = torch.randint(0, n_classes, (n_flip,), generator=g)
# Drop labels
if drop_frac > 0:
n_drop = int(drop_frac * n)
drop_idx = torch.randperm(n, generator=g)[:n_drop]
labels[drop_idx] = -1 # sentinel
# Keep only samples with valid labels
keep = labels >= 0
new_ds.append(WirelessDataset(
name=ds.name,
channels=ds.channels[keep],
labels=labels[keep],
targets=ds.targets[keep] if ds.targets is not None else None,
))
continue
new_ds.append(WirelessDataset(
name=ds.name, channels=ds.channels, labels=labels,
targets=ds.targets,
))
return DatasetCollection(datasets=new_ds)
# ── Run trials ───────────────────────────────────────────────────────────────
all_pearson, all_spearman = [], []
for trial in range(args.trials):
split_seed = args.seed + trial
torch.manual_seed(split_seed)
np.random.seed(split_seed)
# ── Split ────────────────────────────────────────────────────────────
train_ds, val_ds, test_ds = [], [], []
for ds in full:
tr, va, te = ds.split(
train=args.n_train, val=args.n_val, test=args.n_test, seed=split_seed,
)
train_ds.append(tr)
val_ds.append(va)
test_ds.append(te)
train_col = DatasetCollection(datasets=train_ds)
val_col = DatasetCollection(datasets=val_ds)
test_col = DatasetCollection(datasets=test_ds)
# ── Perturb labels for distance computation only ──────────────────
# Keep clean labels for P (ground-truth transfer performance),
# perturb only the copies used for label-aware D.
if args.label_noise > 0 or args.label_drop > 0:
train_col_noisy = _perturb_labels(train_col, args.label_noise, args.label_drop, split_seed)
else:
train_col_noisy = train_col
if trial == 0:
log.info(f"\n{'='*60}")
log.info(f"Generated {len(full)} datasets "
f"(train/val/test = {args.n_train}/{args.n_val}/{args.n_test}):")
for tr, te in zip(train_col, test_col):
n0 = (tr.labels == 0).sum().item()
log.info(f" {tr.name}: train={len(tr)}, test={len(te)}, "
f"LoS%={100*n0/max(len(tr),1):.0f}%")
# ── Load runner ──────────────────────────────────────────────────────
import io, contextlib
with contextlib.redirect_stdout(io.StringIO()):
runner = load_runner(args.lwm_path, weights_file=args.weights_file)
# ── P matrix ─────────────────────────────────────────────────────────
verbose = (args.trials == 1)
P = compute_performance_matrix(
train_col,
cfg=TransferConfig(
task=args.task,
epochs=args.epochs,
lr=args.lr,
batch_size=args.batch_size,
n_beams=args.n_beams,
seed=split_seed,
symmetrize=not (args.asymmetric or args.raw_perf),
verbose=verbose,
),
val_collection=val_col,
test_collection=test_col,
)
# ── CDE (optional) ───────────────────────────────────────────────────
if args.cde:
from dataset_distancing.cde import CDETrainer
from dataset_distancing import CDEConfig
cde_cfg = CDEConfig(
epochs=args.cde_epochs,
lr=args.cde_lr,
n_per_dataset=64,
n_per_dataset_eval=128,
corr_n_per_dataset=64,
asymmetric=args.asymmetric,
distance=DistanceConfig(
sw_projections=args.sw_projections,
label_aware=args.label_aware,
asymmetric=args.asymmetric,
),
)
trainer = CDETrainer(runner, cde_cfg)
trainer.fit(train_col, P, eval_collection=train_col, perf_eval=P, verbose=verbose)
if args.cde_save and trial == 0:
trainer.save(args.cde_save)
# ── D matrix ─────────────────────────────────────────────────────────
if args.embed_space == "lwm":
D = compute_distances(train_col_noisy, runner, cfg=DistanceConfig(
distance=args.distance,
label_aware=args.label_aware,
n_per_dataset=len(train_col[0]),
sw_projections=args.sw_projections,
asymmetric=args.asymmetric,
asymmetric_alpha=args.asym_alpha,
))
else:
# Raw or UMAP: build EmbeddingResult manually
from dataset_distancing.embeddings import EmbeddingResult, channels_to_patches
from dataset_distancing.distances import compute_distances_from_embeddings
# Vectorize channels: (N, 1, ant, sub) complex β†’ (N, 2*ant*sub) float
# Subsample to common size across datasets
n_per = min(len(ds) for ds in train_col_noisy.datasets)
all_embs, all_labels = [], []
for ds in train_col_noisy.datasets:
idx = torch.randperm(len(ds))[:n_per]
ch = ds.channels[idx]
labels = ds.labels[idx] if ds.labels is not None else None
if ch.is_complex():
flat = torch.cat([ch.real.reshape(n_per, -1),
ch.imag.reshape(n_per, -1)], dim=1).float()
else:
flat = ch.reshape(n_per, -1).float()
if args.embed_space == "umap":
import warnings, umap
from sklearn.preprocessing import StandardScaler
with warnings.catch_warnings():
warnings.simplefilter("ignore")
scaled = StandardScaler().fit_transform(flat.numpy())
reducer = umap.UMAP(n_components=args.umap_dim, n_neighbors=16,
min_dist=0.1, random_state=split_seed)
flat = torch.from_numpy(reducer.fit_transform(scaled)).float()
all_embs.append(flat.unsqueeze(0))
all_labels.append(labels)
emb_result = EmbeddingResult(
embeddings=torch.cat(all_embs, dim=0),
labels=all_labels,
dataset_names=train_col_noisy.names(),
)
D = compute_distances_from_embeddings(emb_result, cfg=DistanceConfig(
distance=args.distance,
label_aware=args.label_aware,
sw_projections=args.sw_projections,
asymmetric=args.asymmetric,
asymmetric_alpha=args.asym_alpha,
))
# ── Correlate ────────────────────────────────────────────────────────
stats = correlate(D, P, symmetric=not (args.asymmetric or args.raw_perf))
all_pearson.append(stats['pearson'])
all_spearman.append(stats['spearman'])
if args.trials == 1:
_print_matrix("Distance matrix D", D.numpy())
_print_matrix("Performance matrix P", P.numpy())
log.info(f"\n{'='*60}")
log.info(f"Pearson = {stats['pearson']:.4f}")
log.info(f"Spearman = {stats['spearman']:.4f}")
log.info(f"{'='*60}")
else:
log.info(f" Trial {trial+1}/{args.trials} (seed={split_seed}): "
f"r={stats['pearson']:.4f}, ρ={stats['spearman']:.4f}")
# ── Summary ──────────────────────────────────────────────────────────────────
if args.trials > 1:
p = np.array(all_pearson)
s = np.array(all_spearman)
log.info(f"\n{'='*60}")
log.info(f"Over {args.trials} trials:")
log.info(f" Pearson = {p.mean():.4f} Β± {p.std():.4f}")
log.info(f" Spearman = {s.mean():.4f} Β± {s.std():.4f}")
log.info(f"{'='*60}")
# Examples, from the repository root:
# python tutorial.py --task los-nlos --n_train 400 --n_test 400 --epochs 200 --batch_size 32 --seed 0 --trials 20
# python tutorial.py --task beam-pred --n_beams 32 --n_train 2000 --n_test 2000 --epochs 200 --batch_size 512 --seed 0 --trials 20
# python tutorial.py --task los-nlos --embed_space raw # the same on raw channels
# python tutorial.py --task los-nlos --cde --cde_save ckpt # fine-tune with the LWM-CDE objective first