BattleTag/JANUS
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: code, paper, small artifacts

Browse Source
This commit is contained in:
BattleTag
2026-05-07 20:47:30 +08:00
commit fae2db8cff
322 changed files with 33159 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Mixed_CFM/__init__.py Normal file
View File

@@ -0,0 +1 @@
pass

42
Mixed_CFM/configs/cicddos2019_ac_combo_seed42.yaml Normal file
View File

@@ -0,0 +1,42 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_cicddos2019_seed42
source_store: /home/chy/JANUS/datasets/cicddos2019/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/cicddos2019/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/cicddos2019/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 42
data_seed: 42
train_ratio: 0.8
benign_label: normal
val_cap: 20000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 10000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

42
Mixed_CFM/configs/cicddos2019_ac_combo_seed43.yaml Normal file
View File

@@ -0,0 +1,42 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_cicddos2019_seed43
source_store: /home/chy/JANUS/datasets/cicddos2019/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/cicddos2019/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/cicddos2019/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 43
data_seed: 43
train_ratio: 0.8
benign_label: normal
val_cap: 20000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 10000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

42
Mixed_CFM/configs/cicddos2019_ac_combo_seed44.yaml Normal file
View File

@@ -0,0 +1,42 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_cicddos2019_seed44
source_store: /home/chy/JANUS/datasets/cicddos2019/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/cicddos2019/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/cicddos2019/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 44
data_seed: 44
train_ratio: 0.8
benign_label: normal
val_cap: 20000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 10000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

40
Mixed_CFM/configs/cicids2017_ac_combo_seed42.yaml Normal file
View File

@@ -0,0 +1,40 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_cicids2017_seed42
packets_npz: /home/chy/JANUS/datasets/cicids2017/processed/packets.npz
flows_parquet: /home/chy/JANUS/datasets/cicids2017/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/cicids2017/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 42
data_seed: 42
train_ratio: 0.8
benign_label: normal
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

40
Mixed_CFM/configs/cicids2017_ac_combo_seed43.yaml Normal file
View File

@@ -0,0 +1,40 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_cicids2017_seed43
packets_npz: /home/chy/JANUS/datasets/cicids2017/processed/packets.npz
flows_parquet: /home/chy/JANUS/datasets/cicids2017/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/cicids2017/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 43
data_seed: 43
train_ratio: 0.8
benign_label: normal
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

40
Mixed_CFM/configs/cicids2017_ac_combo_seed44.yaml Normal file
View File

@@ -0,0 +1,40 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_cicids2017_seed44
packets_npz: /home/chy/JANUS/datasets/cicids2017/processed/packets.npz
flows_parquet: /home/chy/JANUS/datasets/cicids2017/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/cicids2017/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 44
data_seed: 44
train_ratio: 0.8
benign_label: normal
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

44
Mixed_CFM/configs/ciciot2023_ac_combo_seed42.yaml Normal file
View File

@@ -0,0 +1,44 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_ciciot2023_seed42
source_store: /home/chy/JANUS/datasets/ciciot2023/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/ciciot2023/processed/full_store/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/ciciot2023/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 42
data_seed: 42
train_ratio: 0.8
benign_label: normal
val_cap: 10000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
device: auto
reference_mode: causal_packets

44
Mixed_CFM/configs/ciciot2023_ac_combo_seed43.yaml Normal file
View File

@@ -0,0 +1,44 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_ciciot2023_seed43
source_store: /home/chy/JANUS/datasets/ciciot2023/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/ciciot2023/processed/full_store/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/ciciot2023/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 43
data_seed: 43
train_ratio: 0.8
benign_label: normal
val_cap: 10000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
device: auto
reference_mode: causal_packets

44
Mixed_CFM/configs/ciciot2023_ac_combo_seed44.yaml Normal file
View File

@@ -0,0 +1,44 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_ciciot2023_seed44
source_store: /home/chy/JANUS/datasets/ciciot2023/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/ciciot2023/processed/full_store/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/ciciot2023/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 44
data_seed: 44
train_ratio: 0.8
benign_label: normal
val_cap: 10000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
device: auto
reference_mode: causal_packets

42
Mixed_CFM/configs/ciciot2023_seed42.yaml Normal file
View File

@@ -0,0 +1,42 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_c_mixed_ciciot2023_seed42
source_store: /home/chy/JANUS/datasets/ciciot2023/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/ciciot2023/processed/full_store/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/ciciot2023/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 42
data_seed: 42
train_ratio: 0.8
benign_label: normal
val_cap: 10000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
device: auto

42
Mixed_CFM/configs/ciciot2023_seed43.yaml Normal file
View File

@@ -0,0 +1,42 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_c_mixed_ciciot2023_seed43
source_store: /home/chy/JANUS/datasets/ciciot2023/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/ciciot2023/processed/full_store/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/ciciot2023/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 43
data_seed: 43
train_ratio: 0.8
benign_label: normal
val_cap: 10000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
device: auto

42
Mixed_CFM/configs/ciciot2023_seed44.yaml Normal file
View File

@@ -0,0 +1,42 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_c_mixed_ciciot2023_seed44
source_store: /home/chy/JANUS/datasets/ciciot2023/processed/full_store
flows_parquet: /home/chy/JANUS/datasets/ciciot2023/processed/full_store/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/ciciot2023/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 44
data_seed: 44
train_ratio: 0.8
benign_label: normal
val_cap: 10000
attack_cap: 20000
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
device: auto

40
Mixed_CFM/configs/iscxtor2016_ac_combo_seed42.yaml Normal file
View File

@@ -0,0 +1,40 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_iscxtor2016_seed42
packets_npz: /home/chy/JANUS/datasets/iscxtor2016/processed/packets.npz
flows_parquet: /home/chy/JANUS/datasets/iscxtor2016/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/iscxtor2016/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 42
data_seed: 42
train_ratio: 0.8
benign_label: nontor
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

40
Mixed_CFM/configs/iscxtor2016_ac_combo_seed43.yaml Normal file
View File

@@ -0,0 +1,40 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_iscxtor2016_seed43
packets_npz: /home/chy/JANUS/datasets/iscxtor2016/processed/packets.npz
flows_parquet: /home/chy/JANUS/datasets/iscxtor2016/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/iscxtor2016/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 43
data_seed: 43
train_ratio: 0.8
benign_label: nontor
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

40
Mixed_CFM/configs/iscxtor2016_ac_combo_seed44.yaml Normal file
View File

@@ -0,0 +1,40 @@
save_dir: /home/chy/JANUS/artifacts/route_comparison/route_ac_combo_iscxtor2016_seed44
packets_npz: /home/chy/JANUS/datasets/iscxtor2016/processed/packets.npz
flows_parquet: /home/chy/JANUS/datasets/iscxtor2016/processed/flows.parquet
flow_features_path: /home/chy/JANUS/datasets/iscxtor2016/processed/flow_features.parquet
flow_features_align: auto
T: 64
n_train: 10000
min_len: 2
seed: 44
data_seed: 44
train_ratio: 0.8
benign_label: nontor
d_model: 128
n_layers: 4
n_heads: 4
mlp_ratio: 4.0
time_dim: 64
token_dim:
batch_size: 256
num_workers: 0
epochs: 50
lr: 3.0e-4
weight_decay: 0.01
grad_clip: 1.0
eval_every: 10
eval_n: 20000
eval_batch_size: 512
eval_n_steps: 8
sigma: 0.1
use_ot: true
lambda_disc: 1.0
reference_mode: causal_packets
device: auto

155
Mixed_CFM/data.py Normal file
View File

@@ -0,0 +1,155 @@
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
import numpy as np
import pandas as pd
import sys as _sys
from pathlib import Path as _Path
_sys.path.insert(0, str(_Path(__file__).resolve().parents[1]))
from common.data_contract import PACKET_FEATURE_NAMES, PACKET_CONTINUOUS_CHANNEL_IDX, PACKET_BINARY_CHANNEL_IDX, fit_packet_stats as _fit_packet_stats, zscore as _zscore
import importlib.util as _ilu
_UDATA_NAME = 'unified_cfm_data'
if _UDATA_NAME not in _sys.modules:
_udata_spec = _ilu.spec_from_file_location(_UDATA_NAME, _Path(__file__).resolve().parents[1] / 'Unified_CFM' / 'data.py')
_udata = _ilu.module_from_spec(_udata_spec)
_sys.modules[_UDATA_NAME] = _udata
_udata_spec.loader.exec_module(_udata)
else:
_udata = _sys.modules[_UDATA_NAME]
DEFAULT_FLOW_META_COLUMNS = _udata.DEFAULT_FLOW_META_COLUMNS
_read_aligned_flow_features = _udata._read_aligned_flow_features
_preprocess_flow = _udata._preprocess_flow
@dataclass
class MixedData:
train_cont: np.ndarray
val_cont: np.ndarray
attack_cont: np.ndarray
train_disc: np.ndarray
val_disc: np.ndarray
attack_disc: np.ndarray
train_flow: np.ndarray
val_flow: np.ndarray
attack_flow: np.ndarray
train_len: np.ndarray
val_len: np.ndarray
attack_len: np.ndarray
attack_labels: np.ndarray
cont_mean: np.ndarray
cont_std: np.ndarray
flow_mean: np.ndarray
flow_std: np.ndarray
flow_feature_names: tuple[str, ...]
packet_feature_names: tuple[str, ...] = PACKET_FEATURE_NAMES
@property
def T(self) -> int:
return int(self.train_cont.shape[1])
@property
def n_cont(self) -> int:
return int(self.train_cont.shape[2])
@property
def n_disc(self) -> int:
return int(self.train_disc.shape[2])
@property
def flow_dim(self) -> int:
return int(self.train_flow.shape[1])
def _zscore_cont(train_x: np.ndarray, val_x: np.ndarray, attack_x: np.ndarray, train_l: np.ndarray, val_l: np.ndarray, attack_l: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
(mean, std) = _fit_packet_stats(train_x, train_l)
def prep(x: np.ndarray, l: np.ndarray) -> np.ndarray:
z = _zscore(x, mean, std)
T = x.shape[1]
m = np.arange(T)[None, :] < l[:, None]
return (z * m[:, :, None]).astype(np.float32)
return (prep(train_x, train_l), prep(val_x, val_l), prep(attack_x, attack_l), mean, std)
def load_mixed_data(*, packets_npz: Path | None=None, source_store: Path | None=None, flows_parquet: Path, flow_features_path: Path, flow_feature_columns: Optional[list[str]]=None, flow_features_align: str='auto', T: int=64, split_seed: int=42, train_ratio: float=0.8, benign_label: str='normal', min_len: int=2, attack_cap: int | None=None, val_cap: int | None=None) -> MixedData:
if (packets_npz is None) == (source_store is None):
raise ValueError('pass exactly one of packets_npz or source_store')
flows_parquet = Path(flows_parquet)
print(f'[data] flows={flows_parquet} packets={(packets_npz if packets_npz else source_store)}')
flow_cols = ['flow_id', 'label', 'src_ip', 'src_port', 'dst_ip', 'dst_port', 'protocol']
flows = pd.read_parquet(flows_parquet, columns=flow_cols)
labels_full = flows['label'].to_numpy().astype(str)
flow_id = flows['flow_id'].to_numpy()
tokens_full: np.ndarray | None = None
store = None
if packets_npz is not None:
pz = np.load(Path(packets_npz))
tokens_full = pz['packet_tokens'].astype(np.float32)
lens_full = pz['packet_lengths'].astype(np.int32)
if T > tokens_full.shape[1]:
raise ValueError(f'requested T={T} > stored {tokens_full.shape[1]}')
tokens_full = tokens_full[:, :T].copy()
lens_full = np.minimum(lens_full, T).astype(np.int32)
if 'flow_id' in pz.files and (not np.array_equal(pz['flow_id'], flow_id)):
raise ValueError('packets_npz / flows_parquet not row-aligned')
else:
from common.packet_store import PacketShardStore
store = PacketShardStore.open(Path(source_store))
store_id = store.read_flows(columns=['flow_id'])['flow_id'].to_numpy()
if not np.array_equal(store_id, flow_id):
raise ValueError('source_store / flows_parquet not row-aligned')
lens_full = np.minimum(store.manifest['packet_length'].to_numpy(dtype=np.int32), T)
(flow_features, flow_names) = _read_aligned_flow_features(Path(flow_features_path), flows, feature_columns=flow_feature_columns, align=flow_features_align)
keep = lens_full >= min_len
labels = labels_full[keep]
flow_features = flow_features[keep]
lens = lens_full[keep]
global_idx = np.flatnonzero(keep).astype(np.int64)
materialized = tokens_full[keep] if tokens_full is not None else None
print(f'[data] kept {keep.sum():,} of {len(keep):,} (min_len={min_len})')
benign = np.where(labels == benign_label)[0]
attack = np.where(labels != benign_label)[0]
rng = np.random.default_rng(split_seed)
rng.shuffle(benign)
n_train = int(len(benign) * train_ratio)
train_local = benign[:n_train]
val_local = benign[n_train:]
if val_cap is not None and len(val_local) > val_cap:
val_local = np.sort(rng.choice(val_local, size=val_cap, replace=False))
if attack_cap is not None and len(attack) > attack_cap:
attack = np.sort(rng.choice(attack, size=attack_cap, replace=False))
print(f'[data] train={len(train_local):,} val={len(val_local):,} attack={len(attack):,}')
def _materialize(idx_local: np.ndarray) -> np.ndarray:
if materialized is not None:
return materialized[idx_local].astype(np.float32, copy=False)
assert store is not None
g = global_idx[idx_local]
(tok, _) = store.read_packets(g.astype(np.int64), T=T)
return tok.astype(np.float32, copy=False)
tr_p = _materialize(train_local)
va_p = _materialize(val_local)
at_p = _materialize(attack)
tr_l = lens[train_local]
va_l = lens[val_local]
at_l = lens[attack]
tr_f = flow_features[train_local]
va_f = flow_features[val_local]
at_f = flow_features[attack]
cont_idx = list(PACKET_CONTINUOUS_CHANNEL_IDX)
disc_idx = list(PACKET_BINARY_CHANNEL_IDX)
tr_cont = tr_p[..., cont_idx]
va_cont = va_p[..., cont_idx]
at_cont = at_p[..., cont_idx]
tr_disc = tr_p[..., disc_idx].astype(np.int8)
va_disc = va_p[..., disc_idx].astype(np.int8)
at_disc = at_p[..., disc_idx].astype(np.int8)
(tr_cont, va_cont, at_cont, c_mean, c_std) = _zscore_cont(tr_cont, va_cont, at_cont, tr_l, va_l, at_l)
(tr_flow, va_flow, at_flow, f_mean, f_std) = _preprocess_flow(tr_f, va_f, at_f)
return MixedData(train_cont=tr_cont, val_cont=va_cont, attack_cont=at_cont, train_disc=tr_disc, val_disc=va_disc, attack_disc=at_disc, train_flow=tr_flow, val_flow=va_flow, attack_flow=at_flow, train_len=tr_l, val_len=va_l, attack_len=at_l, attack_labels=labels[attack], cont_mean=c_mean, cont_std=c_std, flow_mean=f_mean, flow_std=f_std, flow_feature_names=tuple(flow_names))
def subsample_train(data: MixedData, n: int, seed: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
if n <= 0 or n >= len(data.train_cont):
return (data.train_flow, data.train_cont, data.train_disc, data.train_len)
rng = np.random.default_rng(seed)
idx = rng.choice(len(data.train_cont), n, replace=False)
idx.sort()
return (data.train_flow[idx], data.train_cont[idx], data.train_disc[idx], data.train_len[idx])

180
Mixed_CFM/eval_cross.py Normal file
View File

@@ -0,0 +1,180 @@
from __future__ import annotations
import argparse
import json
import sys as _sys
import time
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from sklearn.metrics import average_precision_score, roc_auc_score
REPO = Path(__file__).resolve().parents[1]
_sys.path.insert(0, str(REPO))
from common.data_contract import PACKET_CONTINUOUS_CHANNEL_IDX, PACKET_BINARY_CHANNEL_IDX, zscore as _zscore
from common.packet_store import PacketShardStore
_sys.path.insert(0, str(Path(__file__).resolve().parent))
from model import MixedCFMConfig, MixedTokenCFM
def _device(arg: str) -> torch.device:
if arg == 'auto':
return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
return torch.device(arg)
def _score_batch(model, flow_z, cont_z, disc_int, lens, device, batch_size=256, n_steps=16):
out: dict[str, list[np.ndarray]] = {}
for start in range(0, len(flow_z), batch_size):
sl = slice(start, start + batch_size)
f = torch.from_numpy(flow_z[sl]).float().to(device)
c = torch.from_numpy(cont_z[sl]).float().to(device)
d = torch.from_numpy(disc_int[sl]).long().to(device)
l = torch.from_numpy(lens[sl]).long().to(device)
with torch.no_grad():
traj = model.trajectory_metrics(f, c, d, l, n_steps=n_steps)
nll = model.disc_nll_score(f, c, d, l)
for src in (traj, nll):
for (k, v) in src.items():
out.setdefault(k, []).append(v.detach().cpu().numpy())
if start // batch_size % 20 == 0:
print(f'[score] {min(start + batch_size, len(flow_z)):,}/{len(flow_z):,}', flush=True)
return {k: np.concatenate(v) for (k, v) in out.items()}
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument('--model-dir', type=Path, required=True)
p.add_argument('--target-store', type=Path, default=None, help='Sharded packet store (mutually exclusive with --target-packets-npz).')
p.add_argument('--target-packets-npz', type=Path, default=None, help='Monolithic packets.npz (for datasets without full_store).')
p.add_argument('--target-flows', type=Path, required=True)
p.add_argument('--target-flow-features', type=Path, required=True)
p.add_argument('--out', type=Path, required=True)
p.add_argument('--n-benign', type=int, default=10000)
p.add_argument('--n-attack', type=int, default=10000)
p.add_argument('--benign-label', type=str, default='normal', help="Label string of benign class in target dataset (e.g. 'nontor' for ISCXTor2016).")
p.add_argument('--seed', type=int, default=42)
p.add_argument('--T', type=int, default=64)
p.add_argument('--batch-size', type=int, default=256)
p.add_argument('--n-steps', type=int, default=16)
p.add_argument('--device', type=str, default='auto')
args = p.parse_args()
if (args.target_store is None) == (args.target_packets_npz is None):
p.error('pass exactly one of --target-store or --target-packets-npz')
device = _device(args.device)
ckpt = torch.load(args.model_dir / 'model.pt', map_location='cpu', weights_only=False)
model_cfg = MixedCFMConfig(**ckpt['model_cfg'])
model = MixedTokenCFM(model_cfg).to(device)
model.load_state_dict(ckpt['model_state_dict'])
model.eval()
cont_mean = np.asarray(ckpt['cont_mean'], dtype=np.float32)
cont_std = np.asarray(ckpt['cont_std'], dtype=np.float32)
flow_mean = np.asarray(ckpt['flow_mean'], dtype=np.float32)
flow_std = np.asarray(ckpt['flow_std'], dtype=np.float32)
flow_names = [str(n) for n in ckpt['flow_feature_names']]
print(f'[model] T={args.T} flow_dim={model_cfg.flow_dim}')
flows = pd.read_parquet(args.target_flows, columns=['flow_id', 'label'])
ff = pd.read_parquet(args.target_flow_features)
if not np.array_equal(flows['flow_id'].to_numpy(dtype=np.uint64), ff['flow_id'].to_numpy(dtype=np.uint64)):
raise ValueError('flows and flow_features not row-aligned')
labels = flows['label'].astype(str).to_numpy()
print(f'[data] {len(flows):,} target rows')
rng = np.random.default_rng(args.seed)
benign_idx = np.flatnonzero(labels == args.benign_label)
attack_idx = np.flatnonzero(labels != args.benign_label)
n_benign = min(args.n_benign, len(benign_idx))
if n_benign < args.n_benign:
print(f'[warn] only {len(benign_idx)} benign rows available (asked {args.n_benign})')
b_sel = np.sort(rng.choice(benign_idx, size=n_benign, replace=False))
atk_classes = sorted(set(labels[attack_idx]))
per_class = max(1, args.n_attack // len(atk_classes))
a_sel_chunks = []
for cls in atk_classes:
pool = attack_idx[labels[attack_idx] == cls]
k = min(per_class, len(pool))
if k:
a_sel_chunks.append(rng.choice(pool, size=k, replace=False))
a_sel = np.sort(np.concatenate(a_sel_chunks))
if len(a_sel) > args.n_attack:
a_sel = np.sort(rng.choice(a_sel, size=args.n_attack, replace=False))
print(f'[sample] benign={len(b_sel):,} attack={len(a_sel):,} ({len(atk_classes)} classes)')
cont_idx = list(PACKET_CONTINUOUS_CHANNEL_IDX)
disc_idx = list(PACKET_BINARY_CHANNEL_IDX)
if args.target_store is not None:
store = PacketShardStore.open(args.target_store)
npz_tokens = None
npz_lens = None
else:
store = None
pz = np.load(args.target_packets_npz)
npz_tokens = pz['packet_tokens'][:, :args.T].astype(np.float32)
npz_lens = np.minimum(pz['packet_lengths'], args.T).astype(np.int32)
def _materialize(idx: np.ndarray):
if store is not None:
(tok, l) = store.read_packets(idx, T=args.T)
else:
tok = npz_tokens[idx]
l = npz_lens[idx]
l = np.minimum(l, args.T).astype(np.int32)
f = ff.iloc[idx][flow_names].to_numpy(dtype=np.float64)
f = np.nan_to_num(f, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
return (tok.astype(np.float32), l, f)
print('[read] benign...')
(b_tok, b_len, b_flow) = _materialize(b_sel)
print('[read] attack...')
(a_tok, a_len, a_flow) = _materialize(a_sel)
if cont_mean.shape == (9,):
cm = cont_mean[cont_idx]
cs = cont_std[cont_idx]
else:
cm = cont_mean
cs = cont_std
def _prep(tok, lens):
cont = tok[..., cont_idx]
disc = tok[..., disc_idx].astype(np.int8)
z = _zscore(cont, cm, cs)
m = np.arange(args.T)[None, :] < lens[:, None]
cont_z = (z * m[:, :, None]).astype(np.float32)
return (cont_z, disc)
(b_cont, b_disc) = _prep(b_tok, b_len)
(a_cont, a_disc) = _prep(a_tok, a_len)
b_flow_z = ((b_flow - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
a_flow_z = ((a_flow - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
t0 = time.time()
print('[eval] benign...')
b_scores = _score_batch(model, b_flow_z, b_cont, b_disc, b_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
print(f'[eval] benign done {time.time() - t0:.1f}s')
t0 = time.time()
print('[eval] attack...')
a_scores = _score_batch(model, a_flow_z, a_cont, a_disc, a_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
print(f'[eval] attack done {time.time() - t0:.1f}s')
keys = sorted(set(b_scores) & set(a_scores))
overall = {}
for k in keys:
y = np.r_[np.zeros(len(b_scores[k])), np.ones(len(a_scores[k]))]
s = np.r_[b_scores[k], a_scores[k]]
s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
overall[k] = {'auroc': float(roc_auc_score(y, s)), 'auprc': float(average_precision_score(y, s))}
a_labels = labels[a_sel]
per_cls = {}
for cls in sorted(set(a_labels)):
m = a_labels == cls
per_cls[cls] = {'_n': float(m.sum())}
for k in keys:
y = np.r_[np.zeros(len(b_scores[k])), np.ones(int(m.sum()))]
s = np.r_[b_scores[k], a_scores[k][m]]
s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
try:
per_cls[cls][k] = float(roc_auc_score(y, s))
except ValueError:
per_cls[cls][k] = float('nan')
out = {'model_dir': str(args.model_dir), 'target_store': str(args.target_store), 'n_benign': len(b_sel), 'n_attack': len(a_sel), 'n_score_keys': len(keys), 'overall': overall, 'per_class': per_cls}
args.out.parent.mkdir(parents=True, exist_ok=True)
args.out.write_text(json.dumps(out, indent=2))
npz = args.out.with_suffix('.npz')
save = {'a_labels': a_labels.astype(str)}
for k in keys:
save[f'b_{k}'] = b_scores[k].astype(np.float32)
save[f'a_{k}'] = a_scores[k].astype(np.float32)
np.savez(npz, **save)
print(f'[saved] {args.out}')
if __name__ == '__main__':
main()

109
Mixed_CFM/eval_phase1.py Normal file
View File

@@ -0,0 +1,109 @@
from __future__ import annotations
import argparse
import json
import sys as _sys
import time
from pathlib import Path
from pathlib import Path as _Path
import numpy as np
import torch
import yaml
from sklearn.metrics import average_precision_score, roc_auc_score
_sys.path.insert(0, str(_Path(__file__).resolve().parent))
from data import load_mixed_data
from model import MixedCFMConfig, MixedTokenCFM
def _device(arg: str) -> torch.device:
if arg == 'auto':
return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
return torch.device(arg)
def _score_batch(model: MixedTokenCFM, flow_np: np.ndarray, cont_np: np.ndarray, disc_np: np.ndarray, len_np: np.ndarray, device: torch.device, *, batch_size: int, n_steps: int) -> dict[str, np.ndarray]:
out: dict[str, list[np.ndarray]] = {}
for start in range(0, len(flow_np), batch_size):
sl = slice(start, start + batch_size)
flow = torch.from_numpy(flow_np[sl]).float().to(device)
cont = torch.from_numpy(cont_np[sl]).float().to(device)
disc = torch.from_numpy(disc_np[sl]).long().to(device)
lens = torch.from_numpy(len_np[sl]).long().to(device)
with torch.no_grad():
traj = model.trajectory_metrics(flow, cont, disc, lens, n_steps=n_steps)
nll = model.disc_nll_score(flow, cont, disc, lens)
for d in (traj, nll):
for (k, v) in d.items():
out.setdefault(k, []).append(v.detach().cpu().numpy())
print(f'[score] {min(start + batch_size, len(flow_np)):,}/{len(flow_np):,}', flush=True)
return {k: np.concatenate(v, axis=0) for (k, v) in out.items()}
def _auroc_safe(y, s) -> float:
try:
return float(roc_auc_score(y, s))
except ValueError:
return float('nan')
def _auprc_safe(y, s) -> float:
try:
return float(average_precision_score(y, s))
except ValueError:
return float('nan')
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument('--model-dir', type=Path, required=True)
p.add_argument('--out-dir', type=Path, required=True)
p.add_argument('--n-val-cap', type=int, default=None)
p.add_argument('--n-atk-cap', type=int, default=None)
p.add_argument('--batch-size', type=int, default=256)
p.add_argument('--n-steps', type=int, default=16)
p.add_argument('--device', type=str, default='auto')
args = p.parse_args()
device = _device(args.device)
args.out_dir.mkdir(parents=True, exist_ok=True)
cfg = yaml.safe_load((args.model_dir / 'config.yaml').read_text())
ckpt = torch.load(args.model_dir / 'model.pt', map_location='cpu', weights_only=False)
model_cfg = MixedCFMConfig(**ckpt['model_cfg'])
model = MixedTokenCFM(model_cfg).to(device)
model.load_state_dict(ckpt['model_state_dict'])
model.eval()
print(f'[model] T={model_cfg.T} flow_dim={model_cfg.flow_dim}')
data = load_mixed_data(packets_npz=Path(cfg['packets_npz']) if cfg.get('packets_npz') else None, source_store=Path(cfg['source_store']) if cfg.get('source_store') else None, flows_parquet=Path(cfg['flows_parquet']), flow_features_path=Path(cfg['flow_features_path']), flow_features_align=str(cfg.get('flow_features_align', 'auto')), T=int(cfg['T']), split_seed=int(cfg.get('data_seed', cfg.get('seed', 42))), train_ratio=float(cfg.get('train_ratio', 0.8)), benign_label=str(cfg.get('benign_label', 'normal')), min_len=int(cfg.get('min_len', 2)), attack_cap=int(cfg['attack_cap']) if cfg.get('attack_cap') else None, val_cap=int(cfg['val_cap']) if cfg.get('val_cap') else None)
print(f'[data] val={len(data.val_flow):,} attack={len(data.attack_flow):,}')
rng = np.random.default_rng(0)
(val_flow, val_cont, val_disc, val_len) = (data.val_flow, data.val_cont, data.val_disc, data.val_len)
(atk_flow, atk_cont, atk_disc, atk_len) = (data.attack_flow, data.attack_cont, data.attack_disc, data.attack_len)
atk_labels = data.attack_labels
if args.n_val_cap is not None and len(val_flow) > args.n_val_cap:
idx = np.sort(rng.choice(len(val_flow), size=args.n_val_cap, replace=False))
(val_flow, val_cont, val_disc, val_len) = (val_flow[idx], val_cont[idx], val_disc[idx], val_len[idx])
if args.n_atk_cap is not None and len(atk_flow) > args.n_atk_cap:
idx = np.sort(rng.choice(len(atk_flow), size=args.n_atk_cap, replace=False))
(atk_flow, atk_cont, atk_disc, atk_len) = (atk_flow[idx], atk_cont[idx], atk_disc[idx], atk_len[idx])
atk_labels = atk_labels[idx]
print(f'[eval] scoring val={len(val_flow):,} atk={len(atk_flow):,}')
t0 = time.time()
val = _score_batch(model, val_flow, val_cont, val_disc, val_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
print(f'[eval] val done {time.time() - t0:.1f}s')
t0 = time.time()
atk = _score_batch(model, atk_flow, atk_cont, atk_disc, atk_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
print(f'[eval] atk done {time.time() - t0:.1f}s')
keys = sorted(set(val) & set(atk))
overall: dict[str, dict[str, float]] = {}
for k in keys:
y = np.r_[np.zeros(len(val[k])), np.ones(len(atk[k]))]
s = np.r_[val[k], atk[k]]
s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
overall[k] = {'auroc': _auroc_safe(y, s), 'auprc': _auprc_safe(y, s)}
per_class: dict[str, dict[str, float]] = {}
for c in sorted(set(atk_labels.tolist())):
m = atk_labels == c
per_class[c] = {'_n': float(m.sum())}
for k in keys:
y = np.r_[np.zeros(len(val[k])), np.ones(int(m.sum()))]
s = np.r_[val[k], atk[k][m]]
s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
per_class[c][k] = _auroc_safe(y, s)
np.savez(args.out_dir / 'phase1_scores.npz', val_labels=np.array(['normal'] * len(val_flow)), atk_labels=atk_labels.astype(str), **{f'val_{k}': val[k] for k in keys}, **{f'atk_{k}': atk[k] for k in keys})
json.dump({'overall': overall, 'per_class': per_class}, open(args.out_dir / 'phase1_summary.json', 'w'), indent=2)
print(f'[wrote] {args.out_dir}/phase1_summary.json keys={len(keys)}')
if __name__ == '__main__':
main()

244
Mixed_CFM/model.py Normal file
View File

@@ -0,0 +1,244 @@
from __future__ import annotations
import math
from dataclasses import dataclass, field
import torch
import torch.nn as nn
import torch.nn.functional as F
import importlib.util as _ilu
import sys as _sys
from pathlib import Path as _Path
_UNIFIED_NAME = 'unified_cfm_model'
if _UNIFIED_NAME not in _sys.modules:
_unified_spec = _ilu.spec_from_file_location(_UNIFIED_NAME, _Path(__file__).resolve().parents[1] / 'Unified_CFM' / 'model.py')
_unified = _ilu.module_from_spec(_unified_spec)
_sys.modules[_UNIFIED_NAME] = _unified
_unified_spec.loader.exec_module(_unified)
else:
_unified = _sys.modules[_UNIFIED_NAME]
AdaLNBlock = _unified.AdaLNBlock
SinusoidalTimeEmb = _unified.SinusoidalTimeEmb
_sinkhorn_coupling = _unified._sinkhorn_coupling
@dataclass
class MixedCFMConfig:
T: int = 64
flow_dim: int = 20
n_cont_pkt: int = 3
n_disc_pkt: int = 6
cont_pkt_idx: tuple[int, ...] = (0, 1, 8)
disc_pkt_idx: tuple[int, ...] = (2, 3, 4, 5, 6, 7)
n_disc_classes: int = 2
token_dim: int | None = None
d_model: int = 128
n_layers: int = 4
n_heads: int = 4
mlp_ratio: float = 4.0
time_dim: int = 64
sigma: float = 0.1
use_ot: bool = False
reference_mode: str | None = None
lambda_disc: float = 1.0
disc_path: str = 'uniform'
disc_embed_scale: float = 1.0
def __post_init__(self) -> None:
if len(self.cont_pkt_idx) != self.n_cont_pkt:
raise ValueError('cont_pkt_idx length mismatch n_cont_pkt')
if len(self.disc_pkt_idx) != self.n_disc_pkt:
raise ValueError('disc_pkt_idx length mismatch n_disc_pkt')
if self.disc_path != 'uniform':
raise NotImplementedError(f'disc_path={self.disc_path}')
class MixedVelocity(nn.Module):
def __init__(self, token_dim: int, seq_len: int, n_disc: int, n_classes: int, d_model: int=128, n_layers: int=4, n_heads: int=4, mlp_ratio: float=4.0, time_dim: int=64, reference_mode: str | None=None) -> None:
super().__init__()
if reference_mode not in (None, 'causal_packets', 'causal_all'):
raise ValueError(f'reference_mode={reference_mode!r}')
self.token_dim = token_dim
self.seq_len = seq_len
self.n_disc = n_disc
self.n_classes = n_classes
self.reference_mode = reference_mode
self.input_proj = nn.Linear(token_dim, d_model)
self.pos_emb = nn.Parameter(torch.zeros(1, seq_len, d_model))
self.type_emb = nn.Embedding(2, d_model)
nn.init.trunc_normal_(self.pos_emb, std=0.02)
nn.init.normal_(self.type_emb.weight, std=0.02)
self.time_emb = SinusoidalTimeEmb(time_dim)
self.cond_mlp = nn.Sequential(nn.Linear(time_dim, d_model), nn.SiLU(), nn.Linear(d_model, d_model))
self.blocks = nn.ModuleList([AdaLNBlock(d_model, n_heads, mlp_ratio, cond_dim=d_model) for _ in range(n_layers)])
self.out_norm = nn.LayerNorm(d_model, elementwise_affine=False)
self.head_v = nn.Linear(d_model, token_dim)
self.head_disc = nn.Linear(d_model, n_disc * n_classes)
for layer in (self.head_v, self.head_disc):
nn.init.zeros_(layer.weight)
nn.init.zeros_(layer.bias)
type_ids = torch.ones(seq_len, dtype=torch.long)
type_ids[0] = 0
self.register_buffer('type_ids', type_ids, persistent=False)
def _attn_mask(self, L: int, device: torch.device) -> torch.Tensor | None:
if self.reference_mode is None:
return None
if self.reference_mode == 'causal_packets':
mask = torch.zeros((L, L), dtype=torch.bool, device=device)
if L > 1:
mask[1:, 1:] = torch.triu(torch.ones(L - 1, L - 1, dtype=torch.bool, device=device), diagonal=1)
return mask
return torch.triu(torch.ones(L, L, dtype=torch.bool, device=device), diagonal=1)
def forward(self, x: torch.Tensor, t: torch.Tensor, key_padding_mask: torch.Tensor | None=None) -> tuple[torch.Tensor, torch.Tensor]:
(B, L, _) = x.shape
if t.dim() == 0:
t = t.expand(B)
h = self.input_proj(x)
h = h + self.pos_emb[:, :L, :] + self.type_emb(self.type_ids[:L])[None, :, :]
cond = self.cond_mlp(self.time_emb(t))
attn_mask = self._attn_mask(L, x.device)
for block in self.blocks:
h = block(h, cond, key_padding_mask, attn_mask=attn_mask)
h = self.out_norm(h)
v = self.head_v(h)
d = self.head_disc(h).view(B, L, self.n_disc, self.n_classes)
return (v, d)
class MixedTokenCFM(nn.Module):
def __init__(self, cfg: MixedCFMConfig) -> None:
super().__init__()
self.cfg = cfg
cont_size = cfg.n_cont_pkt + cfg.n_disc_pkt
self.token_dim = cfg.token_dim or 1 + max(cfg.flow_dim, cont_size)
if self.token_dim < 1 + max(cfg.flow_dim, cont_size):
raise ValueError('token_dim too small')
self.seq_len = cfg.T + 1
self.velocity = MixedVelocity(token_dim=self.token_dim, seq_len=self.seq_len, n_disc=cfg.n_disc_pkt, n_classes=cfg.n_disc_classes, d_model=cfg.d_model, n_layers=cfg.n_layers, n_heads=cfg.n_heads, mlp_ratio=cfg.mlp_ratio, time_dim=cfg.time_dim, reference_mode=cfg.reference_mode)
def _embed_disc(self, x_disc_int: torch.Tensor) -> torch.Tensor:
s = self.cfg.disc_embed_scale
return (x_disc_int.float() - 0.5) * s
def build_tokens(self, flow: torch.Tensor, packets_cont: torch.Tensor, x_disc_t_int: torch.Tensor) -> torch.Tensor:
(B, T, Cp) = packets_cont.shape
assert T == self.cfg.T and Cp == self.cfg.n_cont_pkt
z = packets_cont.new_zeros((B, T + 1, self.token_dim))
z[:, 0, 0] = -1.0
z[:, 0, 1:1 + self.cfg.flow_dim] = flow
z[:, 1:, 0] = 1.0
z[:, 1:, 1:1 + self.cfg.n_cont_pkt] = packets_cont
z[:, 1:, 1 + self.cfg.n_cont_pkt:1 + self.cfg.n_cont_pkt + self.cfg.n_disc_pkt] = self._embed_disc(x_disc_t_int)
return z
def key_padding_mask(self, lens: torch.Tensor) -> torch.Tensor:
B = lens.shape[0]
idx = torch.arange(self.cfg.T, device=lens.device)[None, :]
packet_real = idx < lens[:, None]
real = torch.cat([torch.ones(B, 1, dtype=torch.bool, device=lens.device), packet_real], dim=1)
return ~real
def _loss_mask(self, lens: torch.Tensor) -> torch.Tensor:
return (~self.key_padding_mask(lens)).float()
def compute_loss(self, flow: torch.Tensor, packets_cont: torch.Tensor, packets_disc: torch.Tensor, lens: torch.Tensor, *, return_components: bool=False) -> torch.Tensor | dict[str, torch.Tensor]:
(B, T, _) = packets_cont.shape
device = packets_cont.device
mask = self._loss_mask(lens)
kpm = mask == 0
x_1_cont = self.build_tokens(flow, packets_cont, torch.zeros_like(packets_disc))
x_0_cont = torch.randn_like(x_1_cont)
if self.cfg.use_ot:
flat0 = (x_0_cont * mask[:, :, None]).reshape(B, -1)
flat1 = (x_1_cont * mask[:, :, None]).reshape(B, -1)
col = _sinkhorn_coupling(torch.cdist(flat0.float(), flat1.float()))
x_1_cont = x_1_cont[col]
packets_cont = packets_cont[col]
packets_disc = packets_disc[col]
flow = flow[col]
lens = lens[col]
mask = self._loss_mask(lens)
kpm = mask == 0
t = torch.rand(B, device=device)
x_t_cont = (1.0 - t[:, None, None]) * x_0_cont + t[:, None, None] * x_1_cont
if self.cfg.sigma > 0:
std = self.cfg.sigma * torch.sqrt(t * (1.0 - t))[:, None, None]
x_t_cont = x_t_cont + std * torch.randn_like(x_t_cont)
target_cont = x_1_cont - x_0_cont
u = torch.rand(B, T, self.cfg.n_disc_pkt, device=device)
keep = u < t[:, None, None]
rand_disc = torch.randint(0, self.cfg.n_disc_classes, packets_disc.shape, device=device)
x_disc_t = torch.where(keep, packets_disc, rand_disc)
disc_start = 1 + self.cfg.n_cont_pkt
x_t_full = x_t_cont.clone()
x_t_full[:, 1:, disc_start:disc_start + self.cfg.n_disc_pkt] = self._embed_disc(x_disc_t)
(v_pred, d_logits) = self.velocity(x_t_full, t, key_padding_mask=kpm)
v_err = (v_pred - target_cont).square()
v_err[:, :, disc_start:disc_start + self.cfg.n_disc_pkt] = 0.0
v_per_token = v_err.mean(dim=-1)
per_sample = (v_per_token * mask).sum(dim=-1) / mask.sum(dim=-1).clamp_min(1.0)
L_cont = per_sample.mean()
pkt_logits = d_logits[:, 1:]
pkt_real = mask[:, 1:].bool()
corrupt = ~keep & pkt_real[:, :, None]
flat_logits = pkt_logits.reshape(-1, self.cfg.n_disc_classes)
flat_targets = packets_disc.reshape(-1).long()
flat_ce = F.cross_entropy(flat_logits, flat_targets, reduction='none')
flat_ce = flat_ce.view(B, T, self.cfg.n_disc_pkt)
flat_ce = flat_ce * corrupt.float()
denom = corrupt.float().sum().clamp_min(1.0)
L_disc = flat_ce.sum() / denom
total = L_cont + self.cfg.lambda_disc * L_disc
if return_components:
return {'total': total, 'main': L_cont.detach(), 'aux_disc': L_disc.detach(), 'aux_flow': L_cont.new_zeros(()), 'aux_packet': L_cont.new_zeros(())}
return total
@torch.no_grad()
def trajectory_metrics(self, flow: torch.Tensor, packets_cont: torch.Tensor, packets_disc: torch.Tensor, lens: torch.Tensor, n_steps: int=16) -> dict[str, torch.Tensor]:
z = self.build_tokens(flow, packets_cont, packets_disc)
mask = self._loss_mask(lens)
kpm = mask == 0
B = z.shape[0]
dt = 1.0 / n_steps
disc_start = 1 + self.cfg.n_cont_pkt
disc_end = disc_start + self.cfg.n_disc_pkt
disc_embed = z[:, 1:, disc_start:disc_end].clone()
for k in range(n_steps):
t_val = 1.0 - k * dt
t = torch.full((B,), t_val, device=z.device)
(v, _) = self.velocity(z, t, key_padding_mask=kpm)
v[:, :, disc_start:disc_end] = 0.0
z = z - v * dt
z[:, 1:, disc_start:disc_end] = disc_embed
z_real = z * mask[:, :, None]
z_cont = z_real.clone()
z_cont[:, 1:, disc_start:disc_end] = 0.0
packet_count = mask[:, 1:].sum(dim=-1).clamp_min(1.0)
terminal = z_cont.reshape(B, -1).norm(dim=-1) / (mask.sum(dim=-1) * self.token_dim).clamp_min(1.0).sqrt()
terminal_flow = z_cont[:, 0].norm(dim=-1) / math.sqrt(self.token_dim)
terminal_packet = (z_cont[:, 1:] * mask[:, 1:, None]).reshape(B, -1).norm(dim=-1) / (packet_count * self.token_dim).sqrt()
return {'terminal_norm': terminal, 'terminal_flow': terminal_flow, 'terminal_packet': terminal_packet}
@torch.no_grad()
def disc_nll_score(self, flow: torch.Tensor, packets_cont: torch.Tensor, packets_disc: torch.Tensor, lens: torch.Tensor, t_eval: float=0.5) -> dict[str, torch.Tensor]:
(B, T, _) = packets_cont.shape
device = packets_cont.device
mask = self._loss_mask(lens)
kpm = mask == 0
z = self.build_tokens(flow, packets_cont, packets_disc)
t = torch.full((B,), float(t_eval), device=device)
(_, d_logits) = self.velocity(z, t, key_padding_mask=kpm)
pkt_logits = d_logits[:, 1:]
flat_logits = pkt_logits.reshape(-1, self.cfg.n_disc_classes)
flat_targets = packets_disc.reshape(-1).long()
ce = F.cross_entropy(flat_logits, flat_targets, reduction='none')
ce = ce.view(B, T, self.cfg.n_disc_pkt)
pkt_real = mask[:, 1:].bool().float()
per_sample = (ce.sum(dim=-1) * pkt_real).sum(dim=-1) / pkt_real.sum(dim=-1).clamp_min(1.0)
per_ch = (ce * pkt_real[:, :, None]).sum(dim=1) / pkt_real.sum(dim=1).clamp_min(1.0)[:, None]
out = {'disc_nll_total': per_sample}
for (c, idx) in enumerate(self.cfg.disc_pkt_idx):
out[f'disc_nll_ch{idx}'] = per_ch[:, c]
return out
def param_count(self) -> int:
return sum((p.numel() for p in self.parameters()))

141
Mixed_CFM/train.py Normal file
View File

@@ -0,0 +1,141 @@
from __future__ import annotations
import argparse
import json
import sys as _sys
import time
from dataclasses import asdict
from pathlib import Path
from pathlib import Path as _Path
from typing import Any
import numpy as np
import torch
import yaml
from sklearn.metrics import roc_auc_score
from torch.utils.data import DataLoader, TensorDataset
_sys.path.insert(0, str(_Path(__file__).resolve().parent))
from data import MixedData, load_mixed_data, subsample_train
from model import MixedCFMConfig, MixedTokenCFM
def _device(arg: str) -> torch.device:
if arg == 'auto':
return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
return torch.device(arg)
def _batch_score(model: MixedTokenCFM, flow_np: np.ndarray, cont_np: np.ndarray, disc_np: np.ndarray, len_np: np.ndarray, device: torch.device, *, batch_size: int, n_steps: int) -> dict[str, np.ndarray]:
out: dict[str, list[np.ndarray]] = {}
model.eval()
for start in range(0, len(flow_np), batch_size):
sl = slice(start, start + batch_size)
flow = torch.from_numpy(flow_np[sl]).float().to(device)
cont = torch.from_numpy(cont_np[sl]).float().to(device)
disc = torch.from_numpy(disc_np[sl]).long().to(device)
lens = torch.from_numpy(len_np[sl]).long().to(device)
m = model.trajectory_metrics(flow, cont, disc, lens, n_steps=n_steps)
d = model.disc_nll_score(flow, cont, disc, lens)
for src in (m, d):
for (k, v) in src.items():
out.setdefault(k, []).append(v.detach().cpu().numpy())
return {k: np.concatenate(v, axis=0) for (k, v) in out.items()}
def _quick_eval(model: MixedTokenCFM, data: MixedData, device: torch.device, cfg: dict[str, Any]) -> dict[str, float]:
n_eval = int(cfg.get('eval_n', 2000))
rng = np.random.default_rng(0)
def pick(n: int) -> np.ndarray:
m = min(n_eval, n)
return rng.choice(n, m, replace=False)
vi = pick(len(data.val_flow))
ai = pick(len(data.attack_flow))
v = _batch_score(model, data.val_flow[vi], data.val_cont[vi], data.val_disc[vi], data.val_len[vi], device, batch_size=int(cfg.get('eval_batch_size', 512)), n_steps=int(cfg.get('eval_n_steps', 8)))
a = _batch_score(model, data.attack_flow[ai], data.attack_cont[ai], data.attack_disc[ai], data.attack_len[ai], device, batch_size=int(cfg.get('eval_batch_size', 512)), n_steps=int(cfg.get('eval_n_steps', 8)))
y = np.concatenate([np.zeros(len(vi)), np.ones(len(ai))])
out: dict[str, float] = {}
for k in sorted(v.keys()):
s = np.concatenate([v[k], a[k]])
s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
out[f'auroc_{k}'] = float(roc_auc_score(y, s))
return out
def train(cfg: dict[str, Any]) -> Path:
device = _device(str(cfg.get('device', 'auto')))
save_dir = Path(cfg['save_dir'])
save_dir.mkdir(parents=True, exist_ok=True)
with open(save_dir / 'config.yaml', 'w') as f:
yaml.safe_dump(cfg, f)
seed = int(cfg.get('seed', 42))
data_seed = int(cfg.get('data_seed', seed))
torch.manual_seed(seed)
np.random.seed(seed)
print(f'Device: {device} seed=model:{seed}/data:{data_seed}')
data = load_mixed_data(packets_npz=Path(cfg['packets_npz']) if cfg.get('packets_npz') else None, source_store=Path(cfg['source_store']) if cfg.get('source_store') else None, flows_parquet=Path(cfg['flows_parquet']), flow_features_path=Path(cfg['flow_features_path']), flow_feature_columns=cfg.get('flow_feature_columns'), flow_features_align=str(cfg.get('flow_features_align', 'auto')), T=int(cfg['T']), split_seed=data_seed, train_ratio=float(cfg.get('train_ratio', 0.8)), benign_label=str(cfg.get('benign_label', 'normal')), min_len=int(cfg.get('min_len', 2)), attack_cap=int(cfg['attack_cap']) if cfg.get('attack_cap') else None, val_cap=int(cfg['val_cap']) if cfg.get('val_cap') else None)
print(f'[data] T={data.T} cont={data.n_cont} disc={data.n_disc} flow={data.flow_dim} train={len(data.train_flow):,} val={len(data.val_flow):,} attack={len(data.attack_flow):,}')
(tr_f, tr_c, tr_d, tr_l) = subsample_train(data, int(cfg.get('n_train', 0)), data_seed)
ds = TensorDataset(torch.from_numpy(tr_f).float(), torch.from_numpy(tr_c).float(), torch.from_numpy(tr_d).long(), torch.from_numpy(tr_l).long())
loader = DataLoader(ds, batch_size=int(cfg['batch_size']), shuffle=True, drop_last=True, num_workers=int(cfg.get('num_workers', 0)), pin_memory=device.type == 'cuda')
print(f'[data] training on {len(ds):,} flows')
model_cfg = MixedCFMConfig(T=data.T, flow_dim=data.flow_dim, token_dim=cfg.get('token_dim'), d_model=int(cfg['d_model']), n_layers=int(cfg['n_layers']), n_heads=int(cfg['n_heads']), mlp_ratio=float(cfg.get('mlp_ratio', 4.0)), time_dim=int(cfg.get('time_dim', 64)), sigma=float(cfg.get('sigma', 0.1)), use_ot=bool(cfg.get('use_ot', False)), reference_mode=cfg.get('reference_mode'), lambda_disc=float(cfg.get('lambda_disc', 1.0)))
model = MixedTokenCFM(model_cfg).to(device)
print(f'[model] params={model.param_count():,} token_dim={model.token_dim} sigma={model_cfg.sigma} use_ot={model_cfg.use_ot} lambda_disc={model_cfg.lambda_disc}')
opt = torch.optim.AdamW(model.parameters(), lr=float(cfg['lr']), weight_decay=float(cfg.get('weight_decay', 0.01)))
total_steps = max(1, int(cfg['epochs']) * len(loader))
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps)
history: dict[str, list[Any]] = {'epoch': [], 'loss': [], 'eval': []}
for epoch in range(1, int(cfg['epochs']) + 1):
model.train()
losses: list[float] = []
ldisc_sum = 0.0
n_batches = 0
t0 = time.time()
for (flow, cont, disc, lens) in loader:
flow = flow.to(device, non_blocking=True)
cont = cont.to(device, non_blocking=True)
disc = disc.to(device, non_blocking=True)
lens = lens.to(device, non_blocking=True)
comp = model.compute_loss(flow, cont, disc, lens, return_components=True)
loss = comp['total']
ldisc_sum += float(comp['aux_disc'].item())
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), float(cfg.get('grad_clip', 1.0)))
opt.step()
sched.step()
losses.append(float(loss.item()))
n_batches += 1
mean_loss = float(np.mean(losses)) if losses else float('nan')
eval_metrics: dict[str, float] | None = None
if epoch % int(cfg.get('eval_every', 5)) == 0 or epoch == int(cfg['epochs']):
eval_metrics = _quick_eval(model, data, device, cfg)
history['epoch'].append(epoch)
history['loss'].append(mean_loss)
history['eval'].append(eval_metrics)
elapsed = time.time() - t0
tail = ''
if eval_metrics:
t = eval_metrics.get('auroc_terminal_norm', float('nan'))
n = eval_metrics.get('auroc_disc_nll_total', float('nan'))
tail = f' auroc_term={t:.3f} auroc_disc={n:.3f}'
if n_batches:
tail += f' L_disc={ldisc_sum / n_batches:.4f}'
print(f"[epoch {epoch:>3d}/{cfg['epochs']:<3d}] ({elapsed:.1f}s) loss={mean_loss:.4f}{tail}")
if not np.isfinite(mean_loss):
raise RuntimeError(f'non-finite loss at epoch {epoch}')
payload = {'model_state_dict': model.state_dict(), 'model_cfg': asdict(model_cfg), 'cont_mean': data.cont_mean, 'cont_std': data.cont_std, 'flow_mean': data.flow_mean, 'flow_std': data.flow_std, 'flow_feature_names': np.asarray(data.flow_feature_names), 'packet_feature_names': np.asarray(data.packet_feature_names)}
torch.save(payload, save_dir / 'model.pt')
with open(save_dir / 'history.json', 'w') as f:
json.dump(history, f, indent=2, default=str)
print(f"[saved] {save_dir / 'model.pt'}")
return save_dir
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--config', type=Path, required=True)
p.add_argument('--override', type=str, nargs='*', default=[])
args = p.parse_args()
with open(args.config) as f:
cfg = yaml.safe_load(f)
for ov in args.override:
(k, v) = ov.split('=', 1)
cfg[k] = yaml.safe_load(v)
train(cfg)
if __name__ == '__main__':
main()