*** Begin Patch
*** Update File: netflow/model/[Link]
@@
-class RelativeTimeBias([Link]):
- """把相对时间差 Δt 经过 RBF/Fourier basis → 标量偏置，用于注意力 logits。"""
-
- def __init__(self, num_bases: int = 8):
- super().__init__()
- [Link] = [Link]([Link](-2.0, 2.0, num_bases),
requires_grad=False)
- [Link] = [Link]([Link](num_bases), requires_grad=False)
- self.w = [Link]([Link](num_bases))
-
- def forward(self, dt: [Link]) -> [Link]:
- """
- dt: (L, L) 已 log1p 并标准化到 ~[-2,2]
- return: (L, L) 标量偏置
- """
- # RBF φ_k(dt) = exp(-((dt-c_k)^2) / (2σ_k^2))
- bases = [Link](-(([Link](-1) - [Link]) ** 2) / (2 *
([Link] ** 2)))
- bias = [Link](bases, self.w, dims=([-1], [0])) # (L, L)
- return bias
+class RelativeTimeBias([Link]):
+ """
+ 把相对时间差 Δt 经过**可学习的 RBF 基**映射到标量偏置，用于注意力 logits。
+
+ 用法：传入原始 Δt（秒），内部做 log1p 与标准化，再用可学习的 (centers, widths, weights) 生成偏置。
+ """
+
+ def __init__(self, num_bases: int = 8):
+ super().__init__()
+ self.num_bases = num_bases
+ # 初始化到 [-2,2] 等距；设置为可学习参数
+ [Link] = [Link]([Link](-2.0, 2.0, num_bases))
+ self.log_width = [Link]([Link](num_bases)) # width =
exp(log_width)
+ self.w = [Link]([Link](num_bases))
+
+ def forward(self, dt_seconds: [Link]) -> [Link]:
+ """
+ Args:
+ dt_seconds: (L, L) 原始绝对时间差（单位：秒，非负）
+ Returns:
+ bias: (L, L) 标量偏置
+ """
+ # 预处理：log1p + 标准化到 ~[-2,2]
+ x = torch.log1p(dt_seconds)
+ x = (x - [Link]()) / ([Link]() + 1e-6)
+ width = [Link](self.log_width) + 1e-6
+ # RBF φ_k(x) = exp(-((x-c_k)^2)/(2σ_k^2))
+ bases = [Link](-(([Link](-1) - [Link]) ** 2) / (2 * (width
** 2))) # (L,L,K)
+ bias = [Link](bases, self.w, dims=([-1], [0])) # (L,L)
+ # 对角线不加偏置（自身不需要）
+ bias = bias - [Link]([Link](bias))
+ return bias
@@
-class TemporalEncoder([Link]):
+class TemporalEncoder([Link]):
@@
- def __init__(self, d: int, nhead: int = 4, dropout: float = 0.1):
+ def __init__(self, d: int, nhead: int = 4, dropout: float = 0.1, rbf_bases:
int = 8):
super().__init__()
[Link] = PositionalEncoding(d, dropout=dropout)
[Link] = [Link](embed_dim=d, num_heads=nhead,
dropout=dropout, batch_first=False)
self.ln1 = [Link](d)
[Link] = [Link]([Link](d, d * 4), [Link](),
[Link](dropout), [Link](d * 4, d))
self.ln2 = [Link](d)
self.res_scale = [Link]([Link](1e-3))
+ # 可学习相对时间偏置与模态指示偏置
+ [Link] = RelativeTimeBias(num_bases=rbf_bases)
+ self.port_bias = [Link]([Link](0.0))
+ self.dst_bias = [Link]([Link](0.0))
@@
- X: [Link],
- attn_bias: Optional[[Link]] = None,
+ X: [Link],
+ attn_bias: Optional[[Link]] = None,
+ # 下述三个用于**在模型内**计算偏置（推荐）。若提供它们，将覆盖 attn_bias。
+ dt_seconds: Optional[[Link]] = None,
+ same_port: Optional[[Link]] = None,
+ same_dst: Optional[[Link]] = None,
key_padding_mask: Optional[[Link]] = None,
positions: Optional[[Link]] = None,
) -> [Link]:
# X: (L, B, d)
X = [Link](X, positions)
- if attn_bias is not None:
- # MultiheadAttention 的 attn_mask 形状 (L, L)，对 batch 共享。
- # 这里取 batch 内第一条的 bias；或外层循环 per-window 调用。
- attn_mask = attn_bias[0]
- else:
- attn_mask = None
+ # 构造注意力 mask（对 batch 共享，故 per-window 逐次调用）
+ if dt_seconds is not None:
+ mask = [Link](dt_seconds)
+ if same_port is not None:
+ mask = mask + self.port_bias * same_port
+ if same_dst is not None:
+ mask = mask + self.dst_bias * same_dst
+ attn_mask = mask
+ elif attn_bias is not None:
+ attn_mask = attn_bias[0]
+ else:
+ attn_mask = None
H, _ = [Link](X, X, X, attn_mask=attn_mask,
key_padding_mask=key_padding_mask)
Y = self.ln1(X + self.res_scale * H)
Y2 = self.ln2(Y + self.res_scale * [Link](Y))
return Y2
*** End Patch
*** Begin Patch
*** Update File: netflow/data/window_graph.py
@@
-def build_sequence_windows(
+def build_sequence_windows(
batch: List[Dict],
cfg: WindowConfig,
-) -> Tuple[List[List[int]], List[[Link]], [Link]]:
+) -> Tuple[List[List[int]], List[dict], [Link]]:
@@
- seq_bias: List[[Link]] = []
+ seq_bias: List[dict] = []
@@
- # 计算注意力偏置
+ # 计算 pairwise 组件（交由模型内 RBF 形成可学习偏置）
t = times[win_local]
L = len(win_local)
dt = [Link]([Link](L, 1) - [Link](1, L))
- bias = _relative_time_bias(dt)
# 端口/目的地相等偏置
dps = [Link]([dports[j] for j in win_local])
dst_arr = [Link]([dsts[j] for j in win_local], dtype=object)
same_port = ([Link](L, 1) == [Link](1,
L)).astype(np.float32)
same_dst = (dst_arr.reshape(L, 1) == dst_arr.reshape(1,
L)).astype(np.float32)
- bias = bias + cfg.bias_port * same_port + cfg.bias_dst * same_dst
-
- seq_indices.append(win_global)
- seq_bias.append(torch.from_numpy(bias).float())
+ seq_indices.append(win_global)
+ seq_bias.append(
+ {
+ "dt": torch.from_numpy(dt).float(),
+ "same_port": torch.from_numpy(same_port).float(),
+ "same_dst": torch.from_numpy(same_dst).float(),
+ }
+ )
*** End Patch
*** Begin Patch
*** Update File: netflow/model/[Link]
@@
- [Link] = TemporalEncoder(d=d, nhead=4)
+ [Link] = TemporalEncoder(d=d, nhead=4)
@@
- bias = seq_bias[win_id].unsqueeze(0) if seq_bias is not None else None
# (1, L, L)
- Y = [Link](X, attn_bias=bias, key_padding_mask=None) # (L, 1, d)
+ bias_item = seq_bias[win_id] if seq_bias is not None else None
+ if isinstance(bias_item, dict):
+ Y = [Link](
+ X,
+ dt_seconds=bias_item.get("dt").to(device),
+ same_port=bias_item.get("same_port").to(device),
+ same_dst=bias_item.get("same_dst").to(device),
+ key_padding_mask=None,
+ )
+ else:
+ bias = bias_item.unsqueeze(0) if bias_item is not None else None
+ Y = [Link](X, attn_bias=bias, key_padding_mask=None)
*** End Patch
*** Begin Patch
*** Add File: examples/run_minimal_e2e.py
+"""
+最小端到端示例：
+ 读取 CSV → 组装 DataLoader（带窗口/图）→ MultiViewEncoder 前向 → 简单训练一步。
+
+CSV 假定包含列（可用 YAML 改名）：
+ ts, src_ip, dst_ip, sport, dport, app, <数值特征若干列...>
+
+运行：
+ python examples/run_minimal_e2e.py --csv data/[Link] --config
configs/[Link]
+"""
+
+import argparse
+import csv
+from typing import Any, Dict, List
+
+import torch
+from [Link] import Dataset, DataLoader
+
+from [Link] import MultiViewEncoder
+from [Link] import load_yaml, window_config_from_yaml
+from [Link] import default_collate_with_windows
+
+
+class CSVFlowDataset(Dataset):
+ def __init__(self, path: str, fields_map: Dict[str, str], num_cols:
List[str]):
+ [Link]: List[Dict[str, Any]] = []
+ with open(path, "r", newline="") as f:
+ reader = [Link](f)
+ for r in reader:
+ item = {
+ "ts": float(r[fields_map.get("ts", "ts")]),
+ "src_ip": r[fields_map.get("src_ip", "src_ip")],
+ "dst_ip": r[fields_map.get("dst_ip", "dst_ip")],
+ "sport": int([Link](fields_map.get("sport", "sport"), 0) or 0),
+ "dport": int([Link](fields_map.get("dport", "dport"), 0) or 0),
+ "app": [Link](fields_map.get("app", "app"), "") or "",
+ }
+ x_num = [float(r[c]) for c in num_cols]
+ item["x_num"] = [Link](x_num, dtype=torch.float32)
+ [Link](item)
+
+ def __len__(self):
+ return len([Link])
+
+ def __getitem__(self, idx):
+ return [Link][idx]
+
+
+class DummyBase([Link]):
+ def __init__(self, in_dim: int, d: int):
+ super().__init__()
+ [Link] = [Link](
+ [Link](in_dim, d), [Link](), [Link](d, d)
+ )
+
+ def forward(self, x_num, x_cat=None):
+ return [Link](x_num)
+
+
+def main():
+ ap = [Link]()
+ ap.add_argument("--csv", type=str, required=True)
+ ap.add_argument("--config", type=str, default="configs/[Link]")
+ ap.add_argument("--batch", type=int, default=64)
+ ap.add_argument("--d", type=int, default=128)
+ args = ap.parse_args()
+
+ cfg = load_yaml([Link])
+ wcfg = window_config_from_yaml(cfg)
+ fields = [Link]("data", {}).get("fields", {})
+
+ # 推断数值特征列：去掉已知字段
+ known = {[Link](k, k) for k in ["ts", "src_ip", "dst_ip", "sport",
"dport", "app"]}
+ # 从 CSV 头读取一次
+ with open([Link], "r", newline="") as f:
+ header = next([Link](f))
+ num_cols = [c for c in header if c not in known]
+ print("Use numeric cols:", num_cols)
+
+ ds = CSVFlowDataset([Link], fields, num_cols)
+ collate = lambda batch: default_collate_with_windows(batch, wcfg)
+ dl = DataLoader(ds, batch_size=[Link], shuffle=True, collate_fn=collate)
+
+ base = DummyBase(in_dim=len(num_cols), d=args.d)
+ enc = MultiViewEncoder(base_encoder=base, d=args.d, quality_dim=8)
+
+ opt = [Link]([Link](), lr=1e-3)
+
+ [Link]()
+ for step, (x_num, x_cat, ctx, y) in enumerate(dl):
+ if step > 5:
+ break
+ z_flow = enc.encode_flow(x_num)
+ z_flow, z_gnn, z_seq, z_star = enc(
+ z_flow, ctx.seq_indices, ctx.seq_bias, ctx.graph_adj, [Link]
+ )
+ # 演示：用 z* 与 z_flow 的对齐作为简单损失（实际训练请换成你的目标）
+ cos = [Link].cosine_similarity(z_star, z_flow, dim=-1)
+ loss = (1 - cos).mean()
+ opt.zero_grad()
+ [Link]()
+ [Link]()
+ print(f"step {step}: loss={[Link]():.4f}")
+
+ print("Done. You can now plug this into your SupCon/nnPU pipeline.")
+
+
+if __name__ == "__main__":
+ main()
+
*** End Patch
*** Begin Patch
*** Update File: configs/[Link]
@@
temporal:
- use_transformer: true
- nhead: 4
- dropout: 0.1
+ use_transformer: true
+ nhead: 4
+ dropout: 0.1
+ rbf_bases: 8
*** End Patch