930 lines
45 KiB
Python
Executable File
930 lines
45 KiB
Python
Executable File
"""
|
||
=============================================================================
|
||
Multi-User Semantic Communication — PyTorch MAML Training + Evaluation
|
||
IEEE JSAC: Meta-Learned Cross-Attention for Multi-User Semantic
|
||
Communication over Wireless Fading Channels
|
||
|
||
구조
|
||
----
|
||
SemanticEncoder : MLP x_u → e_u ∈ ℝ^d
|
||
UserWiseCrossAttn : user-wise cross-attention decoder Y → ê_u (논문 식 4–6)
|
||
MAMLTrainer : MAML outer/inner loop over SNR tasks (논문 식 7–9)
|
||
|
||
비교 대상
|
||
----------
|
||
OFDMA : 대역폭 B/U 분할 (SNR 패널티 −10log10(U))
|
||
NOMA-SIC : 전력 중첩 + 순차 간섭 제거
|
||
Joint : 고정 SNR 분포에서 표준 joint training
|
||
|
||
지표
|
||
----
|
||
SER : cosine-sim(ê_u, e_u) < τ 인 사용자 비율
|
||
ρ : 디코딩된 임베딩의 사용자 간 Pearson 상관계수
|
||
|
||
실행 방법
|
||
----------
|
||
python maml_semantic.py # 학습 + 평가
|
||
python maml_semantic.py --fast # 빠른 디버그 (epoch 축소)
|
||
python maml_semantic.py --eval_only --ckpt results/models.pt
|
||
=============================================================================
|
||
"""
|
||
|
||
import argparse
|
||
import copy
|
||
import warnings
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
import matplotlib.gridspec as gridspec
|
||
from matplotlib.colors import LinearSegmentedColormap
|
||
|
||
warnings.filterwarnings("ignore")
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 0. CONFIG
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def get_cfg():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--d", type=int, default=64)
|
||
p.add_argument("--U", type=int, default=4)
|
||
p.add_argument("--H", type=int, default=4, help="attention heads")
|
||
p.add_argument("--tau", type=float, default=0.45)
|
||
p.add_argument("--lam", type=float, default=0.1, help="ortho loss weight λ")
|
||
p.add_argument("--snr_min", type=float, default=0.0)
|
||
p.add_argument("--snr_max", type=float, default=20.0)
|
||
p.add_argument("--snr_step", type=float, default=2.0)
|
||
p.add_argument("--inner_lr", type=float, default=0.01)
|
||
p.add_argument("--inner_steps", type=int, default=5)
|
||
p.add_argument("--outer_lr", type=float, default=1e-3)
|
||
p.add_argument("--meta_epochs", type=int, default=300)
|
||
p.add_argument("--joint_epochs",type=int, default=300)
|
||
p.add_argument("--batch", type=int, default=64)
|
||
p.add_argument("--n_mc", type=int, default=500)
|
||
p.add_argument("--seed", type=int, default=42)
|
||
p.add_argument("--outdir", type=str, default="results")
|
||
p.add_argument("--eval_only", action="store_true")
|
||
p.add_argument("--ckpt", type=str, default=None)
|
||
p.add_argument("--fast", action="store_true",
|
||
help="빠른 디버그: epoch을 1/10로 축소")
|
||
p.add_argument("--device", type=str, default="auto")
|
||
p.add_argument("--scenario", type=str, default="DEFAULT",
|
||
help="훈련 시나리오: DEFAULT/HIGH/LOW/MIX/HETERO/ASYM/ALL")
|
||
p.add_argument("--decoder_only", action="store_true",
|
||
help="Option A: IdentityEncoder(고정) + decoder만 MAML 훈련")
|
||
args = p.parse_args()
|
||
if args.fast:
|
||
args.meta_epochs = max(30, args.meta_epochs // 10)
|
||
args.joint_epochs = max(30, args.joint_epochs // 10)
|
||
args.n_mc = max(50, args.n_mc // 10)
|
||
return args
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 1. DATA GENERATION (자율주행 합성 임베딩)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# u0: 보행자 bounding-box (scene 결합도 높음)
|
||
# u1: 신호등 상태 (중간)
|
||
# u2: 차선 세그먼테이션 (낮음)
|
||
# u3: 차량 속도/방향 (가장 낮음)
|
||
BLEND = [0.55, 0.45, 0.30, 0.20]
|
||
USER_LABELS = ['보행자\n(U1)', '신호등\n(U2)', '차선\n(U3)', '속도\n(U4)']
|
||
USER_COLORS = ['#1565C0', '#2E7D32', '#C62828', '#6A1B9A']
|
||
|
||
# Scenario configs matching semantic_correlation_sim.py
|
||
SCENARIO_CONFIGS = {
|
||
'DEFAULT': {'beta_u': [0.55, 0.45, 0.30, 0.20], 'scenes': [0, 0, 0, 0]},
|
||
'HIGH': {'beta_u': [0.65, 0.65, 0.60, 0.60], 'scenes': [0, 0, 0, 0]},
|
||
'LOW': {'beta_u': [0.65, 0.05, 0.05, 0.05], 'scenes': [0, 1, 2, 3]},
|
||
'MIX': {'beta_u': [0.65, 0.65, 0.05, 0.05], 'scenes': [0, 0, 1, 2]},
|
||
'HETERO': {'beta_u': [0.75, 0.75, 0.45, 0.08], 'scenes': [0, 0, 0, 1]},
|
||
'ASYM': {'beta_u': [0.72, 0.58, 0.35, 0.12], 'scenes': [0, 0, 0, 0]},
|
||
}
|
||
|
||
|
||
def gen_embeddings(n: int, d: int, U: int, rng,
|
||
scenario_cfg=None) -> torch.Tensor:
|
||
"""단위 정규화된 ground-truth 임베딩 (n, U, d)
|
||
scenario_cfg: dict with 'beta_u' and 'scenes' keys (from SCENARIO_CONFIGS).
|
||
If None, uses the default BLEND with a single shared scene.
|
||
"""
|
||
if scenario_cfg is None:
|
||
blend = BLEND
|
||
scenes = [0] * U
|
||
else:
|
||
blend = scenario_cfg['beta_u']
|
||
scenes = scenario_cfg['scenes']
|
||
|
||
# One unit-norm scene vector per unique scene key (full D dims, matches sim)
|
||
unique_scenes = sorted(set(scenes))
|
||
scene_vecs: dict = {}
|
||
for sc in unique_scenes:
|
||
v = rng.standard_normal(d)
|
||
scene_vecs[sc] = v / (np.linalg.norm(v) + 1e-8)
|
||
|
||
embs = []
|
||
for u in range(U):
|
||
b = blend[u] if u < len(blend) else 0.15
|
||
private = rng.standard_normal((n, d))
|
||
p_hat = private / (np.linalg.norm(private, axis=-1, keepdims=True) + 1e-8)
|
||
s = scene_vecs[scenes[u]]
|
||
# Eq. (2): e_u = sqrt(1-β²)·p̂_u + β·s (semantic_correlation_sim.py convention)
|
||
e = np.sqrt(max(1 - b ** 2, 0)) * p_hat + b * s[None, :]
|
||
e /= np.linalg.norm(e, axis=-1, keepdims=True) + 1e-8
|
||
embs.append(e)
|
||
return torch.from_numpy(np.stack(embs, axis=1)).float() # (n, U, d)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 2. CHANNEL MODELS (SE — Shared Embedding superposition framework)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def _block_masks(U: int, d: int, device) -> torch.Tensor:
|
||
"""Hard block masks (U, d): user u owns dims [u·DPU, (u+1)·DPU).
|
||
Matches semantic_correlation_sim.py MASKS construction exactly."""
|
||
DPU = d // U
|
||
masks = torch.zeros(U, d, device=device)
|
||
for u in range(U):
|
||
masks[u, u * DPU:(u + 1) * DPU] = 1.0
|
||
return masks
|
||
|
||
|
||
def se_channel(E: torch.Tensor, snr_db: float) -> torch.Tensor:
|
||
"""SE superposition channel (논문 Eq. 3).
|
||
|
||
x_u = e_u ⊙ m_u (hard block mask; user u owns D/U contiguous dims)
|
||
y_tx = Σ_u x_u (superimposed D-dim signal)
|
||
y_rx,u = h_u · y_tx + n_u (independent Rayleigh per user)
|
||
|
||
E : (n, U, d) unit-norm embeddings
|
||
Returns Y_rx : (n, U, d) — row u is user u's received copy of y_tx.
|
||
|
||
Cross-attention 구조 설명:
|
||
R_{u,i} = y_rx,u ⊙ m_i 는 user u의 수신 신호에서 subspace i를 추출.
|
||
h_u · x_i[dims_i] + noise 만 남으므로, beta_ui > 0이면 e_u 정보를 담고 있음.
|
||
"""
|
||
n, U, d = E.shape
|
||
masks = _block_masks(U, d, E.device) # (U, d)
|
||
X = E * masks[None, :, :] # (n, U, d)
|
||
Ytx = X.sum(1) # (n, d) superimposed
|
||
|
||
h = (torch.randn(n, U, 1, device=E.device) ** 2 +
|
||
torch.randn(n, U, 1, device=E.device) ** 2).sqrt() * (0.5 ** 0.5) # (n, U, 1)
|
||
|
||
snr_lin = 10 ** (snr_db / 10)
|
||
noise_std = (Ytx.pow(2).mean() / snr_lin).sqrt()
|
||
Yrx = h * Ytx[:, None, :] + torch.randn(n, U, d, device=E.device) * noise_std
|
||
return Yrx # (n, U, d)
|
||
|
||
|
||
def se_ofdma_decode(Y_rx: torch.Tensor) -> torch.Tensor:
|
||
"""SE-OFDMA decoder: ê_u = normalize(y_rx,u ⊙ m_u).
|
||
Each user uses only their own D/U-dim subspace block; no SNR penalty.
|
||
Matches ofdma_se_decoder() in semantic_correlation_sim.py."""
|
||
n, U, d = Y_rx.shape
|
||
masks = _block_masks(U, d, Y_rx.device) # (U, d)
|
||
return F.normalize(Y_rx * masks[None, :, :], dim=-1) # (n, U, d)
|
||
|
||
|
||
def se_noma_decode(E: torch.Tensor, snr_db: float) -> torch.Tensor:
|
||
"""Full-band power-domain NOMA with SIC (standard NOMA, fair comparison).
|
||
|
||
All U users share the SAME full d-dim band — no subspace masking — so the
|
||
received signal is a single d-dim power-domain superposition:
|
||
x_u = sqrt(p_u) · e_u (full-band, power-weighted)
|
||
y = Σ_u h_u · sqrt(p_u) · e_u + n (single d-dim signal at BS)
|
||
|
||
This reflects NOMA's defining resource advantage fairly: every user accesses
|
||
the full d-dim band (vs. OFDMA's exclusive d/U-dim block) at the cost of
|
||
inter-user interference, while the total transmit power Σ_u p_u = 1 matches
|
||
the OFDMA budget, keeping the comparison both power- and bandwidth-fair.
|
||
Each user is reconstructed from the full d-dim signal via successive
|
||
interference cancellation in descending allocated-power order.
|
||
|
||
Matches noma_ul_channel()/noma_sic_decoder() in semantic_correlation_sim.py
|
||
and noma_ch()/noma_sic() in revision_realdata_plot.py.
|
||
|
||
E : (n, U, d) unit-norm ground-truth embeddings.
|
||
Returns decoded : (n, U, d) per-user estimates over the full d-dim band.
|
||
"""
|
||
n, U, d = E.shape
|
||
pa = torch.tensor([0.40, 0.30, 0.20, 0.10], device=E.device)[:U]
|
||
pa = pa / pa.sum() # Σ p_u = 1 (power-fair)
|
||
|
||
h = (torch.randn(n, U, 1, device=E.device) ** 2 +
|
||
torch.randn(n, U, 1, device=E.device) ** 2).sqrt() * (0.5 ** 0.5) # (n, U, 1)
|
||
|
||
y = (E * pa.sqrt()[None, :, None] * h).sum(1) # (n, d) full-band superposition
|
||
snr_lin = 10 ** (snr_db / 10)
|
||
noise_std = (y.pow(2).mean() / snr_lin).sqrt()
|
||
y = y + torch.randn(n, d, device=E.device) * noise_std # (n, d)
|
||
|
||
order = torch.argsort(pa, descending=True) # high-power first
|
||
res = y.clone() # (n, d) shared residual
|
||
decoded = torch.zeros(n, U, d, device=E.device) # (n, U, d)
|
||
for u_idx in order:
|
||
u = int(u_idx.item())
|
||
ê_u = F.normalize(res / (h[:, u, :] + 1e-8), dim=-1) # (n, d) over full band
|
||
decoded[:, u, :] = ê_u
|
||
res = res - h[:, u, :] * pa[u].sqrt() * ê_u # full-band cancellation
|
||
|
||
return decoded # (n, U, d)
|
||
|
||
|
||
def se_sfdma_decode(E: torch.Tensor, snr_db: float) -> torch.Tensor:
|
||
"""SFDMA — Semantic Feature Division Multiple Access (Ma et al., 2024).
|
||
|
||
Each user maps its embedding onto an assigned orthonormal d/U-dim *semantic*
|
||
subspace; all users transmit simultaneously over the SAME full d-dim band, and
|
||
the receiver separates them by projecting onto each user's subspace:
|
||
x_u = P_u e_u (P_u: projector onto user u's semantic subspace)
|
||
y = Σ_u h_u · x_u + n (full-band superposition)
|
||
ê_u = normalize(P_u y / h_u)
|
||
|
||
Fair vs. NOMA/UWCA: SFDMA shares the full band (not OFDMA's exclusive 1/U
|
||
physical sub-band) and is given the best case of *perfectly* orthogonal
|
||
subspaces, so the projections separate users with no inter-user interference.
|
||
The subspaces are taken as the canonical orthogonal partition (the block
|
||
basis) so the comparison isolates the multiple-access mechanism rather than an
|
||
arbitrary basis–data alignment. Because the orthogonality confines each user
|
||
to d/U effective dimensions — the same subspace ceiling as OFDMA — and discards
|
||
the inter-user correlation, SFDMA coincides with OFDMA despite using the full
|
||
band. Matches sfdma() in revision_realdata_plot.py.
|
||
|
||
E : (n, U, d) unit-norm embeddings. Returns (n, U, d) estimates.
|
||
"""
|
||
n, U, d = E.shape
|
||
masks = _block_masks(U, d, E.device) # (U, d) orthogonal subspaces
|
||
|
||
h = (torch.randn(n, U, 1, device=E.device) ** 2 +
|
||
torch.randn(n, U, 1, device=E.device) ** 2).sqrt() * (0.5 ** 0.5) # (n, U, 1)
|
||
|
||
X = E * masks[None, :, :] # (n, U, d) project onto own subspace
|
||
y = (h * X).sum(1) # (n, d) full-band superposition
|
||
snr_lin = 10 ** (snr_db / 10)
|
||
noise_std = (y.pow(2).mean() / snr_lin).sqrt()
|
||
y = y + torch.randn(n, d, device=E.device) * noise_std # (n, d)
|
||
|
||
return torch.stack(
|
||
[F.normalize((y * masks[u]) / (h[:, u, :] + 1e-8), dim=-1) for u in range(U)], 1)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 3. NEURAL MODULES
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
class SemanticEncoder(nn.Module):
|
||
"""f_φ : ℝ^d → ℝ^d (MLP + LayerNorm + 단위 정규화)"""
|
||
def __init__(self, d: int):
|
||
super().__init__()
|
||
self.net = nn.Sequential(
|
||
nn.Linear(d, d * 2), nn.LayerNorm(d * 2), nn.GELU(),
|
||
nn.Linear(d * 2, d),
|
||
)
|
||
|
||
def forward(self, x):
|
||
return F.normalize(self.net(x), dim=-1)
|
||
|
||
|
||
class IdentityEncoder(nn.Module):
|
||
"""No-op encoder: L2-normalizes input only (no learnable params).
|
||
Used for decoder-only training (Option A) so that the channel model
|
||
directly transmits the raw embeddings, matching the analytical simulation."""
|
||
def forward(self, x):
|
||
return F.normalize(x, dim=-1)
|
||
|
||
|
||
class UserWiseCrossAttention(nn.Module):
|
||
"""
|
||
User-wise cross-attention decoder (논문 Section III-A, 식 4–6)
|
||
|
||
입력 Y : (n, U, d) 수신 임베딩 (잡음 포함)
|
||
출력 Ê : (n, U, d) 정제된 임베딩
|
||
복잡도 O(U d²)
|
||
"""
|
||
def __init__(self, d: int, U: int, H: int = 4):
|
||
super().__init__()
|
||
assert d % H == 0
|
||
self.d = d; self.U = U; self.H = H; self.dk = d // H
|
||
# JSAC convention: learnable per-user query vectors {q_u}, not signal-derived
|
||
self.q_vectors = nn.Parameter(torch.randn(U, d) * (d ** -0.5)) # (U, d)
|
||
self.eta = nn.Parameter(torch.ones(1)) # sharpness η
|
||
self.W_K = nn.Linear(d, d, bias=False)
|
||
self.W_V = nn.Linear(d, d, bias=False)
|
||
self.W_O = nn.Linear(d, d, bias=False)
|
||
self.norm = nn.LayerNorm(d)
|
||
# 학습 가능한 soft mask m_i ∈ [0,1]^d (식 4)
|
||
# 초기화: hard block mask logit (+3 = sigmoid → 0.95, -3 → 0.05)
|
||
# 이 초기화가 없으면 모든 mask = 0.5 → K feature가 모든 i에서 동일
|
||
# → softmax가 항상 uniform(1/U) → 학습 무력화
|
||
init_logits = torch.full((U, d), -3.0)
|
||
DPU = d // U
|
||
for u_ in range(U):
|
||
init_logits[u_, u_ * DPU:(u_ + 1) * DPU] = 3.0 # own block 강조
|
||
self.mask_logits = nn.Parameter(init_logits)
|
||
|
||
def forward(self, Y):
|
||
"""Y : (n, U, d) → (Ê, alpha)"""
|
||
n, U, d = Y.shape
|
||
masks = torch.sigmoid(self.mask_logits) # (U, d)
|
||
|
||
# R_{u,i} = y_u ⊙ m_i → (n, U_q, U_k, d)
|
||
R = Y.unsqueeze(2) * masks[None, None, :, :] # broadcast
|
||
|
||
# JSAC: q_u is a fixed learnable vector per user, independent of received signal
|
||
Q = self.q_vectors.unsqueeze(0).expand(n, -1, -1) # (n, U, d)
|
||
K = self.W_K(R) # (n, U, U, d)
|
||
V = self.W_V(R)
|
||
|
||
# multi-head reshape
|
||
def mh(t):
|
||
return t.reshape(*t.shape[:-1], self.H, self.dk)
|
||
|
||
Q_ = mh(Q) # (n,Uq,H,dk)
|
||
K_ = mh(K) # (n,Uq,Uk,H,dk)
|
||
V_ = mh(V)
|
||
|
||
# score : q_u^T k_i / sqrt(dk)
|
||
Q_e = Q_.unsqueeze(3) # (n,Uq,H,1,dk)
|
||
K_t = K_.permute(0, 1, 3, 2, 4) # (n,Uq,H,Uk,dk)
|
||
scores = self.eta * (Q_e * K_t).sum(-1) / (self.dk ** 0.5) # η * q_u^T k_i / √dk
|
||
alpha = F.softmax(scores, dim=-1) # (n,Uq,H,Uk)
|
||
|
||
V_t = V_.permute(0, 1, 3, 2, 4) # (n,Uq,H,Uk,dk)
|
||
ctx = (alpha.unsqueeze(-1) * V_t).sum(3) # (n,Uq,H,dk)
|
||
ctx = self.W_O(ctx.reshape(n, U, d)) # (n,Uq,d)
|
||
|
||
own = Y * masks[None, :, :] # 자기 유저 residual
|
||
out = F.normalize(self.norm(ctx + own), dim=-1)
|
||
return out, alpha.mean(0) # alpha: (Uq,H,Uk)
|
||
|
||
|
||
class SemanticCommSystem(nn.Module):
|
||
def __init__(self, d, U, H, decoder_only=False):
|
||
super().__init__()
|
||
self.encoder = IdentityEncoder() if decoder_only else SemanticEncoder(d)
|
||
self.decoder = UserWiseCrossAttention(d, U, H)
|
||
|
||
def forward(self, X, snr_db):
|
||
n, U, d = X.shape
|
||
E = self.encoder(X.reshape(n * U, d)).reshape(n, U, d)
|
||
Y = se_channel(E, snr_db) # always SE superposition channel
|
||
Ehat, alpha = self.decoder(Y)
|
||
return E, Ehat, alpha
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 4. LOSS (논문 식 8)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def semantic_loss(Ehat, E, lam=0.1):
|
||
"""
|
||
L = (1/U)Σ(1 - cos(ê_u,e_u)) + λ Σ_{u≠v}|ρ(ê_u,ê_v)|
|
||
"""
|
||
cos = (Ehat * E).sum(-1)
|
||
distortion = (1 - cos).mean()
|
||
|
||
U = Ehat.shape[1]
|
||
emb = Ehat.mean(0) # (U, d)
|
||
ec = emb - emb.mean(1, keepdim=True)
|
||
en = F.normalize(ec, dim=1)
|
||
C = en @ en.T
|
||
mask = ~torch.eye(U, dtype=torch.bool, device=Ehat.device)
|
||
ortho = C[mask].abs().mean()
|
||
|
||
return distortion + lam * ortho, distortion.item(), ortho.item()
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 5. MAML TRAINER (논문 식 7, 9)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
class MAMLTrainer:
|
||
def __init__(self, model, cfg, device, rng, scenario_cfg=None):
|
||
self.model = model
|
||
self.cfg = cfg
|
||
self.device = device
|
||
self.rng = rng
|
||
self.scenario_cfg = scenario_cfg # None → default BLEND
|
||
mask_params = [p for n, p in model.named_parameters() if 'mask_logits' in n]
|
||
other_params = [p for n, p in model.named_parameters() if 'mask_logits' not in n]
|
||
self.optimizer = torch.optim.Adam([
|
||
{'params': other_params, 'lr': cfg.outer_lr},
|
||
{'params': mask_params, 'lr': cfg.outer_lr * 100},
|
||
])
|
||
self.snr_tasks = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
|
||
|
||
def _inner_adapt(self, snr):
|
||
"""Inner-loop: adapted copy of model for task T_k = SNR γ_k"""
|
||
adapted = copy.deepcopy(self.model)
|
||
opt_in = torch.optim.SGD(adapted.parameters(), lr=self.cfg.inner_lr)
|
||
for _ in range(self.cfg.inner_steps):
|
||
X = gen_embeddings(self.cfg.batch, self.cfg.d,
|
||
self.cfg.U, self.rng,
|
||
self.scenario_cfg).to(self.device)
|
||
E, Ehat, _ = adapted(X, snr)
|
||
loss, _, _ = semantic_loss(Ehat, E, self.cfg.lam)
|
||
opt_in.zero_grad(); loss.backward(); opt_in.step()
|
||
return adapted
|
||
|
||
def meta_step(self):
|
||
"""Outer-loop: first-order MAML (Reptile-style aggregation)"""
|
||
fo_loss = torch.tensor(0.0, device=self.device)
|
||
for snr in self.snr_tasks:
|
||
X = gen_embeddings(self.cfg.batch, self.cfg.d,
|
||
self.cfg.U, self.rng,
|
||
self.scenario_cfg).to(self.device)
|
||
E, Ehat, _ = self.model(X, float(snr))
|
||
loss, _, _ = semantic_loss(Ehat, E, self.cfg.lam)
|
||
fo_loss = fo_loss + loss
|
||
fo_loss = fo_loss / len(self.snr_tasks)
|
||
self.optimizer.zero_grad()
|
||
fo_loss.backward()
|
||
nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
|
||
self.optimizer.step()
|
||
return fo_loss.item()
|
||
|
||
def train(self):
|
||
history = []
|
||
print("=" * 60)
|
||
print(f"MAML Training ({self.cfg.meta_epochs} epochs, "
|
||
f"{len(self.snr_tasks)} SNR tasks)")
|
||
print("=" * 60)
|
||
for ep in range(1, self.cfg.meta_epochs + 1):
|
||
loss = self.meta_step()
|
||
history.append(loss)
|
||
if ep % max(1, self.cfg.meta_epochs // 6) == 0:
|
||
print(f" Epoch {ep:4d}/{self.cfg.meta_epochs} loss={loss:.4f}")
|
||
return history
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 6. JOINT TRAINING BASELINE
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def train_joint(model, cfg, device, rng, scenario_cfg=None):
|
||
mask_params = [p for n, p in model.named_parameters() if 'mask_logits' in n]
|
||
other_params = [p for n, p in model.named_parameters() if 'mask_logits' not in n]
|
||
opt = torch.optim.Adam([
|
||
{'params': other_params, 'lr': cfg.outer_lr},
|
||
{'params': mask_params, 'lr': cfg.outer_lr * 100},
|
||
])
|
||
history = []
|
||
print("=" * 60)
|
||
print(f"Joint Training ({cfg.joint_epochs} epochs)")
|
||
print("=" * 60)
|
||
for ep in range(1, cfg.joint_epochs + 1):
|
||
snr = float(np.random.uniform(cfg.snr_min, cfg.snr_max))
|
||
X = gen_embeddings(cfg.batch, cfg.d, cfg.U, rng, scenario_cfg).to(device)
|
||
E, Ehat, _ = model(X, snr)
|
||
loss, _, _ = semantic_loss(Ehat, E, cfg.lam)
|
||
opt.zero_grad(); loss.backward()
|
||
nn.utils.clip_grad_norm_(model.parameters(), 5.0)
|
||
opt.step()
|
||
history.append(loss.item())
|
||
if ep % max(1, cfg.joint_epochs // 6) == 0:
|
||
print(f" Epoch {ep:4d}/{cfg.joint_epochs} loss={loss.item():.4f}")
|
||
return history
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 7. EVALUATION (semantic_sim.py 구조 그대로 유지)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def _normalize_np(E):
|
||
return E / (np.linalg.norm(E, axis=-1, keepdims=True) + 1e-8)
|
||
|
||
def cos_mean(Eh, Egt):
|
||
return (Eh * Egt).sum(-1).mean()
|
||
|
||
def ser_total(Eh, Egt, tau):
|
||
return ((Eh * Egt).sum(-1) < tau).mean()
|
||
|
||
def ser_per_user(Eh, Egt, tau):
|
||
return ((Eh * Egt).sum(-1) < tau).mean(0) # (U,)
|
||
|
||
def corr_matrix(Eh):
|
||
e = Eh.mean(0) # (U, D)
|
||
ec = e - e.mean(1, keepdims=True)
|
||
en = ec / (np.linalg.norm(ec, axis=1, keepdims=True) + 1e-8)
|
||
return en @ en.T # (U, U)
|
||
|
||
|
||
@torch.no_grad()
|
||
def evaluate_model(model, cfg, device, rng, channel="se", scenario_cfg=None):
|
||
"""
|
||
SNR_DB 전 구간에 걸쳐 SER / cosine / per-user SER 계산.
|
||
semantic_sim.py의 시뮬레이션 루프와 동일한 방식.
|
||
"""
|
||
SNR_DB = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
|
||
model.eval()
|
||
res = {"ser": [], "cos": [], "sp": [], "rho": []}
|
||
|
||
for snr in SNR_DB:
|
||
ser_acc = 0.0; cos_acc = 0.0
|
||
sp_acc = np.zeros(cfg.U); rho_acc = []
|
||
|
||
for _ in range(cfg.n_mc):
|
||
Egt = gen_embeddings(cfg.batch, cfg.d, cfg.U, rng, scenario_cfg)
|
||
X = Egt.to(device)
|
||
|
||
# 채널 통과 — 모든 방식이 SE 수퍼포지션 채널 공유
|
||
if channel == "ofdma":
|
||
# SE-OFDMA: 같은 SE 채널, own D/U block만 사용
|
||
E = F.normalize(X.reshape(-1, cfg.d), dim=-1).reshape(-1, cfg.U, cfg.d)
|
||
Y = se_channel(E, float(snr))
|
||
Ehat = se_ofdma_decode(Y)
|
||
elif channel == "noma":
|
||
# SE-NOMA-SIC: full-band power-domain superposition (own channel + SIC)
|
||
E = F.normalize(X.reshape(-1, cfg.d), dim=-1).reshape(-1, cfg.U, cfg.d)
|
||
Ehat = se_noma_decode(E, float(snr))
|
||
elif channel == "sfdma":
|
||
# SFDMA: full-band, orthogonal semantic-subspace division (Ma 2024)
|
||
E = F.normalize(X.reshape(-1, cfg.d), dim=-1).reshape(-1, cfg.U, cfg.d)
|
||
Ehat = se_sfdma_decode(E, float(snr))
|
||
else:
|
||
# MAML+Attn / Joint+Attn: SE 채널 + cross-attention decoder
|
||
E, Ehat, _ = model(X, float(snr))
|
||
|
||
Eh = Ehat.cpu().numpy()
|
||
Egt_np = E.cpu().numpy()
|
||
|
||
ser_acc += ser_total(Eh, Egt_np, cfg.tau)
|
||
cos_acc += cos_mean(Eh, Egt_np)
|
||
sp_acc += ser_per_user(Eh, Egt_np, cfg.tau)
|
||
rho_acc.append(corr_matrix(Eh))
|
||
|
||
res["ser"].append(ser_acc / cfg.n_mc)
|
||
res["cos"].append(cos_acc / cfg.n_mc)
|
||
res["sp"].append(sp_acc / cfg.n_mc)
|
||
res["rho"].append(np.mean(rho_acc, axis=0))
|
||
|
||
for k in res:
|
||
res[k] = np.array(res[k])
|
||
return res
|
||
|
||
|
||
@torch.no_grad()
|
||
def get_attention_map(model, snr, cfg, device, rng):
|
||
model.eval()
|
||
attn_sum = None
|
||
for _ in range(cfg.n_mc):
|
||
X = gen_embeddings(cfg.batch, cfg.d, cfg.U, rng).to(device)
|
||
E = model.encoder(X.reshape(-1, cfg.d)).reshape(-1, cfg.U, cfg.d)
|
||
Y = se_channel(E, float(snr))
|
||
_, alpha = model.decoder(Y) # (Uq, H, Uk)
|
||
a = alpha.mean(1).cpu().numpy() # (Uq, Uk) — avg over heads
|
||
attn_sum = a if attn_sum is None else attn_sum + a
|
||
return attn_sum / cfg.n_mc
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 8. 9-PANEL FIGURE (semantic_sim.py 그대로 재현)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def plot_results(SNR_DB, res_maml, res_joint, res_ofdma, res_noma,
|
||
rho_j, rho_m, attn_m_mc, out_path, cfg):
|
||
|
||
U = cfg.UOF
|
||
KP = 'MAML+Attn\n(제안)'
|
||
MCFG = {
|
||
'OFDMA': ('#546E7A', 's--', 1.6, 'OFDMA'),
|
||
'NOMA-SIC': ('#E65100', '^-.', 1.6, 'NOMA-SIC'),
|
||
'Joint+Attn': ('#C62828', 'D--', 1.8, 'Joint+Attn'),
|
||
KP: ('#1565C0', 'o-', 2.5, 'MAML+Attn (제안)'),
|
||
}
|
||
res_all = {'OFDMA': res_ofdma, 'NOMA-SIC': res_noma,
|
||
'Joint+Attn': res_joint, KP: res_maml}
|
||
|
||
cmap_r = LinearSegmentedColormap.from_list(
|
||
'r', ['#1565C0', '#FFFFFF', '#C62828'], N=256)
|
||
cmap_a = LinearSegmentedColormap.from_list(
|
||
'a', ['#F5F5F5', '#1565C0'], N=256)
|
||
|
||
fig = plt.figure(figsize=(18, 15))
|
||
fig.patch.set_facecolor('#F8F9FA')
|
||
gs = gridspec.GridSpec(3, 3, figure=fig, hspace=0.48, wspace=0.38,
|
||
left=0.07, right=0.97, top=0.93, bottom=0.06)
|
||
|
||
idx10 = int((10 - cfg.snr_min) / cfg.snr_step) # SNR=10dB 인덱스
|
||
|
||
# (a) SER vs SNR
|
||
ax = fig.add_subplot(gs[0, 0]); ax.set_facecolor('white')
|
||
for k, (c, mk, lw, lb) in MCFG.items():
|
||
ax.semilogy(SNR_DB, res_all[k]['ser'], mk, lw=lw, ms=6, color=c, label=lb)
|
||
d10 = res_ofdma['ser'][idx10] - res_maml['ser'][idx10]
|
||
ax.annotate(f'Δ={d10:.3f}\n@ 10 dB',
|
||
xy=(10, res_maml['ser'][idx10]),
|
||
xytext=(13, res_maml['ser'][idx10] * 4),
|
||
fontsize=8.5, color='#1565C0',
|
||
arrowprops=dict(arrowstyle='->', color='#1565C0', lw=1.2))
|
||
ax.set_xlabel('SNR (dB)', fontsize=11); ax.set_ylabel('SER', fontsize=11)
|
||
ax.set_title('(a) SER vs SNR', fontsize=12, fontweight='bold')
|
||
ax.legend(fontsize=9); ax.grid(True, alpha=0.35); ax.set_xlim(0, 20)
|
||
|
||
# (b) 코사인 유사도
|
||
ax = fig.add_subplot(gs[0, 1]); ax.set_facecolor('white')
|
||
for k, (c, mk, lw, lb) in MCFG.items():
|
||
ax.plot(SNR_DB, res_all[k]['cos'], mk, lw=lw, ms=6, color=c, label=lb)
|
||
ax.axhline(cfg.tau, color='gray', lw=1.2, ls=':', label=f'τ={cfg.tau}')
|
||
ax.set_xlabel('SNR (dB)', fontsize=11); ax.set_ylabel('코사인 유사도', fontsize=11)
|
||
ax.set_title('(b) 코사인 유사도 vs SNR', fontsize=12, fontweight='bold')
|
||
ax.legend(fontsize=9); ax.grid(True, alpha=0.35)
|
||
ax.set_xlim(0, 20); ax.set_ylim(0.35, 1.02)
|
||
|
||
# (c) SER 개선량
|
||
ax = fig.add_subplot(gs[0, 2]); ax.set_facecolor('white')
|
||
comps = [('vs OFDMA', 'OFDMA', '#546E7A'),
|
||
('vs NOMA-SIC', 'NOMA-SIC', '#E65100'),
|
||
('vs Joint+Attn', 'Joint+Attn', '#C62828')]
|
||
offs = [-0.3, 0.0, 0.3]
|
||
for (lb, base, col), off in zip(comps, offs):
|
||
ax.bar(SNR_DB + off, res_all[base]['ser'] - res_maml['ser'],
|
||
width=0.28, alpha=0.80, color=col, label=lb)
|
||
ax.axhline(0, color='black', lw=0.8)
|
||
ax.set_xlabel('SNR (dB)', fontsize=11); ax.set_ylabel('SER 개선량', fontsize=11)
|
||
ax.set_title('(c) SER 개선량 (베이스라인 − 제안)', fontsize=12, fontweight='bold')
|
||
ax.legend(fontsize=9); ax.grid(True, alpha=0.25, axis='y')
|
||
|
||
# (d) 제안 사용자별 SER
|
||
ax = fig.add_subplot(gs[1, 0]); ax.set_facecolor('white')
|
||
for ui in range(U):
|
||
ax.semilogy(SNR_DB, res_maml['sp'][:, ui], 'o-', lw=1.8, ms=5,
|
||
color=USER_COLORS[ui], label=USER_LABELS[ui])
|
||
ax.set_xlabel('SNR (dB)', fontsize=11); ax.set_ylabel('SER', fontsize=11)
|
||
ax.set_title('(d) 제안 — 사용자별 SER', fontsize=12, fontweight='bold')
|
||
ax.legend(fontsize=8); ax.grid(True, alpha=0.35); ax.set_xlim(0, 20)
|
||
|
||
# (e) OFDMA 사용자별 SER
|
||
ax = fig.add_subplot(gs[1, 1]); ax.set_facecolor('white')
|
||
for ui in range(U):
|
||
ax.semilogy(SNR_DB, res_ofdma['sp'][:, ui], 's--', lw=1.6, ms=5,
|
||
color=USER_COLORS[ui], label=USER_LABELS[ui])
|
||
ax.set_xlabel('SNR (dB)', fontsize=11); ax.set_ylabel('SER', fontsize=11)
|
||
ax.set_title('(e) OFDMA — 사용자별 SER', fontsize=12, fontweight='bold')
|
||
ax.legend(fontsize=8); ax.grid(True, alpha=0.35); ax.set_xlim(0, 20)
|
||
|
||
# (f) SER @ 10 dB 막대
|
||
ax = fig.add_subplot(gs[1, 2]); ax.set_facecolor('white')
|
||
ms = ['OFDMA', 'NOMA-SIC', 'Joint+Attn', KP]
|
||
s10 = [res_all[m]['ser'][idx10] for m in ms]
|
||
lb10 = ['OFDMA', 'NOMA-SIC', 'Joint\n+Attn', 'MAML+Attn\n(제안)']
|
||
c10 = ['#546E7A', '#E65100', '#C62828', '#1565C0']
|
||
bars = ax.bar(range(4), s10, color=c10, width=0.55,
|
||
edgecolor='white', linewidth=1.2)
|
||
ax.set_xticks(range(4)); ax.set_xticklabels(lb10, fontsize=9.5)
|
||
ax.set_ylabel('SER @ 10 dB', fontsize=11)
|
||
ax.set_title('(f) 방법별 SER @ 10 dB', fontsize=12, fontweight='bold')
|
||
ax.grid(True, alpha=0.3, axis='y')
|
||
for b, v, c in zip(bars, s10, c10):
|
||
ax.text(b.get_x() + b.get_width() / 2, v + 0.003, f'{v:.3f}',
|
||
ha='center', va='bottom', fontsize=10, fontweight='bold', color=c)
|
||
|
||
# (g) 상관계수 — Joint
|
||
ax = fig.add_subplot(gs[2, 0]); ax.set_facecolor('white')
|
||
im = ax.imshow(rho_j, cmap=cmap_r, vmin=-0.3, vmax=0.3, aspect='auto')
|
||
ax.set_xticks(range(U)); ax.set_yticks(range(U))
|
||
ax.set_xticklabels([f'U{i+1}' for i in range(U)], fontsize=10)
|
||
ax.set_yticklabels([f'U{i+1}' for i in range(U)], fontsize=10)
|
||
for i in range(U):
|
||
for j in range(U):
|
||
v = rho_j[i, j]
|
||
ax.text(j, i, f'{v:.3f}', ha='center', va='center', fontsize=11,
|
||
fontweight='bold', color='white' if abs(v) > 0.15 else 'black')
|
||
plt.colorbar(im, ax=ax, fraction=0.046)
|
||
ax.set_title('(g) 상관계수 — Joint Training', fontsize=12, fontweight='bold')
|
||
ax.set_xlabel('사용자 j', fontsize=10); ax.set_ylabel('사용자 i', fontsize=10)
|
||
|
||
# (h) 상관계수 — MAML
|
||
ax = fig.add_subplot(gs[2, 1]); ax.set_facecolor('white')
|
||
im = ax.imshow(rho_m, cmap=cmap_r, vmin=-0.3, vmax=0.3, aspect='auto')
|
||
ax.set_xticks(range(U)); ax.set_yticks(range(U))
|
||
ax.set_xticklabels([f'U{i+1}' for i in range(U)], fontsize=10)
|
||
ax.set_yticklabels([f'U{i+1}' for i in range(U)], fontsize=10)
|
||
for i in range(U):
|
||
for j in range(U):
|
||
v = rho_m[i, j]
|
||
ax.text(j, i, f'{v:.3f}', ha='center', va='center', fontsize=11,
|
||
fontweight='bold', color='white' if abs(v) > 0.15 else 'black')
|
||
plt.colorbar(im, ax=ax, fraction=0.046)
|
||
ax.set_title('(h) 상관계수 — MAML (제안)', fontsize=12, fontweight='bold')
|
||
ax.set_xlabel('사용자 j', fontsize=10); ax.set_ylabel('사용자 i', fontsize=10)
|
||
|
||
# (i) Attention heatmap
|
||
ax = fig.add_subplot(gs[2, 2]); ax.set_facecolor('white')
|
||
im = ax.imshow(attn_m_mc, cmap=cmap_a,
|
||
vmin=0, vmax=attn_m_mc.max(), aspect='auto')
|
||
sh = ['보행자\n(U1)', '신호등\n(U2)', '차선\n(U3)', '속도\n(U4)']
|
||
ax.set_xticks(range(U)); ax.set_yticks(range(U))
|
||
ax.set_xticklabels(sh[:U], fontsize=9); ax.set_yticklabels(sh[:U], fontsize=9)
|
||
for i in range(U):
|
||
for j in range(U):
|
||
v = attn_m_mc[i, j]
|
||
ax.text(j, i, f'{v:.3f}', ha='center', va='center', fontsize=11,
|
||
fontweight='bold',
|
||
color='white' if v > attn_m_mc.max() * 0.5 else '#0D1B3E')
|
||
plt.colorbar(im, ax=ax, fraction=0.046)
|
||
ax.set_title(f'(i) 어텐션 가중치 α_{{u,i}} — MAML @ 10dB',
|
||
fontsize=12, fontweight='bold')
|
||
ax.set_xlabel('참조 사용자 i', fontsize=10)
|
||
ax.set_ylabel('질의 사용자 u', fontsize=10)
|
||
|
||
fig.suptitle(
|
||
'Multi-User Semantic Communication: User-Wise Attention vs Orthogonal Allocation\n'
|
||
f'(자율주행 시나리오 — U={U}, d={cfg.d}, Rayleigh Fading)',
|
||
fontsize=13, fontweight='bold', y=0.97)
|
||
|
||
plt.savefig(out_path, dpi=150, bbox_inches='tight', facecolor='#F8F9FA')
|
||
plt.close()
|
||
print(f" 그림 저장 → {out_path}")
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 9. SUMMARY PRINT (semantic_sim.py 스타일)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def print_summary(res_maml, res_joint, res_ofdma, res_noma,
|
||
rho_j, rho_m, attn_m, cfg):
|
||
SNR_DB = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
|
||
idx = {4: int((4 - cfg.snr_min) / cfg.snr_step),
|
||
10: int((10 - cfg.snr_min) / cfg.snr_step),
|
||
16: int((16 - cfg.snr_min) / cfg.snr_step)}
|
||
U = cfg.U
|
||
mask = ~np.eye(U, dtype=bool)
|
||
|
||
print("\n" + "=" * 62)
|
||
print("NUMERICAL SUMMARY")
|
||
print("=" * 62)
|
||
print(f"{'Method':<22}{'SER@4dB':>9}{'SER@10dB':>10}"
|
||
f"{'SER@16dB':>10}{'Cos@10dB':>10}")
|
||
print("-" * 62)
|
||
for lb, res in [('OFDMA', res_ofdma), ('NOMA-SIC', res_noma),
|
||
('Joint+Attn', res_joint), ('MAML+Attn (제안)', res_maml)]:
|
||
print(f"{lb:<22}{res['ser'][idx[4]]:>9.4f}"
|
||
f"{res['ser'][idx[10]]:>10.4f}"
|
||
f"{res['ser'][idx[16]]:>10.4f}"
|
||
f"{res['cos'][idx[10]]:>10.4f}")
|
||
|
||
print(f"\n임베딩 상관계수 |ρ| (off-diag @ 10 dB):")
|
||
print(f" Joint : mean={np.abs(rho_j[mask]).mean():.4f} "
|
||
f"[{rho_j[mask].min():.4f}, {rho_j[mask].max():.4f}]")
|
||
print(f" MAML : mean={np.abs(rho_m[mask]).mean():.4f} "
|
||
f"[{rho_m[mask].min():.4f}, {rho_m[mask].max():.4f}]")
|
||
|
||
print(f"\n어텐션 가중치 α (MAML @ 10 dB):")
|
||
hdr = ''.join([f" U{j+1}" for j in range(U)])
|
||
print(f"{'':>14}{hdr}")
|
||
names = ['보행자', '신호등', '차선 ', '속도 ']
|
||
for i in range(U):
|
||
row = ''.join([f" {attn_m[i,j]:>7.4f}" for j in range(U)])
|
||
nm = names[i] if i < len(names) else f'U{i+1} '
|
||
print(f" U{i+1}({nm}){row}")
|
||
|
||
print(f"\n핵심: α[보행자→신호등]={attn_m[0,1]:.4f} (높음) vs "
|
||
f"α[보행자→속도]={attn_m[0,3]:.4f} (낮음)")
|
||
print(f" SER 개선 vs OFDMA @ 10dB: "
|
||
f"{res_ofdma['ser'][idx[10]] - res_maml['ser'][idx[10]]:.4f}")
|
||
print("=" * 62)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 10. JSON EXPORT (for overlay in semantic_correlation_sim.py)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
import json
|
||
|
||
def export_results_json(scenario: str, SNR_DB, res_maml, res_joint,
|
||
res_ofdma, res_noma, outdir: Path):
|
||
"""Export SER results to JSON so semantic_correlation_sim.py can overlay them."""
|
||
data = {
|
||
"scenario": scenario,
|
||
"decoder_only": True,
|
||
"snr_db": SNR_DB.tolist(),
|
||
"maml_ser": res_maml["ser"].tolist(),
|
||
"joint_ser": res_joint["ser"].tolist(),
|
||
"ofdma_ser": res_ofdma["ser"].tolist(),
|
||
"noma_ser": res_noma["ser"].tolist(),
|
||
}
|
||
path = outdir / f"trained_{scenario}.json"
|
||
with open(path, "w") as f:
|
||
json.dump(data, f, indent=2)
|
||
print(f"결과 JSON 저장 → {path}")
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 11. MAIN
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
def _run_one_scenario(scenario: str, cfg, device, np_rng, outdir: Path):
|
||
"""Train + evaluate one scenario; save checkpoint + JSON."""
|
||
scenario_cfg = SCENARIO_CONFIGS.get(scenario) # None for DEFAULT falls back to BLEND
|
||
|
||
def make_model():
|
||
return SemanticCommSystem(cfg.d, cfg.U, cfg.H,
|
||
decoder_only=cfg.decoder_only).to(device)
|
||
|
||
maml_model = make_model()
|
||
joint_model = make_model()
|
||
|
||
SNR_DB = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
|
||
|
||
# ── 학습 ──────────────────────────────────────────────────────────
|
||
tag = f"[{scenario}]"
|
||
print(f"\n{'='*60}\n{tag} decoder_only={cfg.decoder_only}\n{'='*60}")
|
||
|
||
trainer = MAMLTrainer(maml_model, cfg, device, np_rng, scenario_cfg)
|
||
hist_m = trainer.train()
|
||
hist_j = train_joint(joint_model, cfg, device, np_rng, scenario_cfg)
|
||
|
||
ckpt_path = outdir / f"models_{scenario}.pt"
|
||
torch.save({"maml": maml_model.state_dict(),
|
||
"joint": joint_model.state_dict()}, ckpt_path)
|
||
print(f"체크포인트 저장 → {ckpt_path}")
|
||
|
||
# 학습 곡선
|
||
fig, ax = plt.subplots(figsize=(7, 4))
|
||
ax.plot(hist_m, label='MAML outer loss', color='#1565C0', lw=1.5)
|
||
ax.plot(hist_j, label='Joint loss', color='#C62828', lw=1.5, ls='--')
|
||
ax.set_xlabel('Epoch'); ax.set_ylabel('Loss')
|
||
ax.set_title(f'Training Loss — {scenario}')
|
||
ax.legend(); ax.grid(True, alpha=0.35); fig.tight_layout()
|
||
fig.savefig(str(outdir / f"training_curves_{scenario}.png"), dpi=150)
|
||
plt.close()
|
||
|
||
# ── 평가 ──────────────────────────────────────────────────────────
|
||
print(f"{tag} 평가 중 ...")
|
||
res_maml = evaluate_model(maml_model, cfg, device, np_rng, "rayleigh", scenario_cfg)
|
||
res_joint = evaluate_model(joint_model, cfg, device, np_rng, "rayleigh", scenario_cfg)
|
||
res_ofdma = evaluate_model(joint_model, cfg, device, np_rng, "ofdma", scenario_cfg)
|
||
res_noma = evaluate_model(joint_model, cfg, device, np_rng, "noma", scenario_cfg)
|
||
|
||
export_results_json(scenario, SNR_DB, res_maml, res_joint,
|
||
res_ofdma, res_noma, outdir)
|
||
|
||
# SNR=10dB 상관계수 & attention map
|
||
idx10 = int((10 - cfg.snr_min) / cfg.snr_step)
|
||
rho_j = res_joint["rho"][idx10]
|
||
rho_m = res_maml["rho"][idx10]
|
||
attn_m = get_attention_map(maml_model, 10.0, cfg, device, np_rng)
|
||
|
||
print_summary(res_maml, res_joint, res_ofdma, res_noma,
|
||
rho_j, rho_m, attn_m, cfg)
|
||
|
||
# 단일 시나리오 그림 저장 (DEFAULT는 기존 filename 유지)
|
||
suffix = "" if scenario == "DEFAULT" else f"_{scenario}"
|
||
fig_path = str(outdir / f"semantic_results{suffix}.png")
|
||
plot_results(SNR_DB, res_maml, res_joint, res_ofdma, res_noma,
|
||
rho_j, rho_m, attn_m, fig_path, cfg)
|
||
return res_maml, res_joint, res_ofdma, res_noma
|
||
|
||
|
||
def main():
|
||
cfg = get_cfg()
|
||
torch.manual_seed(cfg.seed)
|
||
np_rng = np.random.default_rng(cfg.seed)
|
||
|
||
if cfg.device == "auto":
|
||
device = torch.device(
|
||
"cuda" if torch.cuda.is_available() else
|
||
"mps" if torch.backends.mps.is_available() else
|
||
"cpu")
|
||
else:
|
||
device = torch.device(cfg.device)
|
||
print(f"Device: {device} | d={cfg.d}, U={cfg.U}, "
|
||
f"meta_epochs={cfg.meta_epochs}, n_mc={cfg.n_mc}, "
|
||
f"scenario={cfg.scenario}, decoder_only={cfg.decoder_only}")
|
||
|
||
outdir = Path(cfg.outdir)
|
||
outdir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# eval_only mode (legacy)
|
||
if cfg.eval_only and cfg.ckpt:
|
||
SNR_DB = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
|
||
scenario_cfg = SCENARIO_CONFIGS.get(cfg.scenario)
|
||
maml_model = SemanticCommSystem(cfg.d, cfg.U, cfg.H,
|
||
decoder_only=cfg.decoder_only).to(device)
|
||
joint_model = SemanticCommSystem(cfg.d, cfg.U, cfg.H,
|
||
decoder_only=cfg.decoder_only).to(device)
|
||
ck = torch.load(cfg.ckpt, map_location=device)
|
||
maml_model.load_state_dict(ck["maml"])
|
||
joint_model.load_state_dict(ck["joint"])
|
||
print(f"체크포인트 로드 ← {cfg.ckpt}")
|
||
res_maml = evaluate_model(maml_model, cfg, device, np_rng, "rayleigh", scenario_cfg)
|
||
res_joint = evaluate_model(joint_model, cfg, device, np_rng, "rayleigh", scenario_cfg)
|
||
res_ofdma = evaluate_model(joint_model, cfg, device, np_rng, "ofdma", scenario_cfg)
|
||
res_noma = evaluate_model(joint_model, cfg, device, np_rng, "noma", scenario_cfg)
|
||
export_results_json(cfg.scenario, SNR_DB, res_maml, res_joint,
|
||
res_ofdma, res_noma, outdir)
|
||
print(f"\n완료! → {outdir}/")
|
||
return
|
||
|
||
# Train / eval scenarios
|
||
if cfg.scenario.upper() == "ALL":
|
||
scenarios = ['HIGH', 'LOW', 'MIX', 'HETERO', 'ASYM']
|
||
else:
|
||
scenarios = [cfg.scenario]
|
||
|
||
for sc in scenarios:
|
||
_run_one_scenario(sc, cfg, device, np_rng, outdir)
|
||
|
||
print(f"\n완료! → {outdir}/")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|