| |
| """ |
| train_slice.py — Train a GNN classifier on backward program-slice graphs. |
| |
| §11 experiment: does concentrating the GNN on the backward data-flow slice |
| from dangerous sinks (strcpy, memcpy, malloc, free, GEP-with-variable-index) |
| improve over the full-function instruction-level classifier (§7: 58.00%)? |
| |
| Hypothesis: a 15-50 node slice where every node is on the dependency path to |
| a dangerous operation is easier to classify than a 400-node full function graph |
| where 0.5-1% of nodes carry vulnerability signal. |
| |
| Architecture: identical to train_instr.py (RGCN + AttentionalAggregation). |
| The sliced graphs have the same format — same opcode vocab, same 3 edge types. |
| |
| Usage: |
| python train_slice.py --epochs 30 --hidden 64 # baseline comparison |
| python train_slice.py --epochs 60 --hidden 128 # extended run |
| """ |
|
|
| import argparse |
| import pickle |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch_geometric.data import Data |
| from torch_geometric.loader import DataLoader |
| from torch_geometric.nn import RGCNConv |
| from torch_geometric.nn.aggr import AttentionalAggregation |
|
|
| HERE = Path(__file__).parent |
| DATA = HERE / "data" |
|
|
| VOCAB_SIZE = 110 |
|
|
|
|
| |
| |
| |
|
|
| class SliceGNN(nn.Module): |
| def __init__(self, vocab: int = VOCAB_SIZE, embed_dim: int = 128, hidden: int = 64): |
| super().__init__() |
| self.embed = nn.Embedding(vocab, embed_dim, padding_idx=79) |
| self.conv1 = RGCNConv(embed_dim, hidden, num_relations=3) |
| self.conv2 = RGCNConv(hidden, hidden, num_relations=3) |
| gate_nn = nn.Sequential( |
| nn.Linear(hidden, hidden // 2), |
| nn.ReLU(), |
| nn.Dropout(0.3), |
| nn.Linear(hidden // 2, 1), |
| ) |
| self.pool = AttentionalAggregation(gate_nn=gate_nn) |
| self.lin = nn.Linear(hidden, 1) |
|
|
| def forward(self, x, edge_index, edge_type, batch): |
| h = self.embed(x.squeeze(-1)) |
| h = F.relu(self.conv1(h, edge_index, edge_type)) |
| h = F.dropout(h, p=0.3, training=self.training) |
| h = F.relu(self.conv2(h, edge_index, edge_type)) |
| h = self.pool(h, batch) |
| return self.lin(h).squeeze(-1) |
|
|
|
|
| |
| |
| |
|
|
| def load_graphs(pkl_path: Path) -> list[Data]: |
| with open(pkl_path, "rb") as f: |
| raw = pickle.load(f) |
| dataset = [] |
| for g in raw: |
| x = torch.tensor(g["x"], dtype=torch.long) |
| edge_index = torch.tensor(g["edge_index"], dtype=torch.long) |
| edge_type = torch.tensor(g["edge_type"], dtype=torch.long) |
| y = torch.tensor([g["y"]], dtype=torch.float) |
| dataset.append(Data(x=x, edge_index=edge_index, |
| edge_type=edge_type, y=y)) |
| return dataset |
|
|
|
|
| |
| |
| |
|
|
| def train_epoch(model, loader, optimizer, device, pos_weight=None): |
| model.train() |
| total_loss = 0.0 |
| for batch in loader: |
| batch = batch.to(device) |
| optimizer.zero_grad() |
| logits = model(batch.x, batch.edge_index, batch.edge_type, batch.batch) |
| loss = F.binary_cross_entropy_with_logits( |
| logits, batch.y.squeeze(), pos_weight=pos_weight) |
| loss.backward() |
| optimizer.step() |
| total_loss += loss.item() * batch.num_graphs |
| return total_loss / len(loader.dataset) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, loader, device): |
| model.eval() |
| correct = total = 0 |
| for batch in loader: |
| batch = batch.to(device) |
| logits = model(batch.x, batch.edge_index, batch.edge_type, batch.batch) |
| preds = (logits > 0).long() |
| labels = batch.y.squeeze().long() |
| correct += (preds == labels).sum().item() |
| total += batch.num_graphs |
| return correct / total if total > 0 else 0.0 |
|
|
|
|
| |
| |
| |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--epochs", type=int, default=30) |
| ap.add_argument("--hidden", type=int, default=64) |
| ap.add_argument("--embed-dim", type=int, default=128) |
| ap.add_argument("--lr", type=float, default=1e-3) |
| ap.add_argument("--batch-size", type=int, default=32) |
| ap.add_argument("--checkpoint", type=str, default="model_slice.pt") |
| args = ap.parse_args() |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"\nDevice: {device}") |
|
|
| for split in ["train", "valid", "test"]: |
| if not (DATA / f"{split}_slice_graphs.pkl").exists(): |
| print(f"Missing data/{split}_slice_graphs.pkl -- run preprocess_slice.py first.") |
| sys.exit(1) |
|
|
| print("Loading slice graphs ...") |
| train_data = load_graphs(DATA / "train_slice_graphs.pkl") |
| valid_data = load_graphs(DATA / "valid_slice_graphs.pkl") |
| test_data = load_graphs(DATA / "test_slice_graphs.pkl") |
| print(f" train={len(train_data)} valid={len(valid_data)} test={len(test_data)}") |
|
|
| node_counts = [d.x.shape[0] for d in train_data] |
| import statistics |
| print(f" Train slice sizes: mean={statistics.mean(node_counts):.0f} " |
| f"median={statistics.median(node_counts):.0f} max={max(node_counts)}") |
|
|
| vuln_train = sum(1 for d in train_data if d.y.item() == 1) |
| fixed_train = len(train_data) - vuln_train |
| pos_weight = torch.tensor([fixed_train / vuln_train]).to(device) |
| print(f" train class balance: {vuln_train} vuln / {fixed_train} fixed") |
| print(f" pos_weight: {pos_weight.item():.3f}\n") |
|
|
| train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True) |
| valid_loader = DataLoader(valid_data, batch_size=args.batch_size) |
| test_loader = DataLoader(test_data, batch_size=args.batch_size) |
|
|
| model = SliceGNN(VOCAB_SIZE, args.embed_dim, args.hidden).to(device) |
| optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) |
| scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) |
|
|
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"Model: SliceGNN(vocab={VOCAB_SIZE}, embed={args.embed_dim}, " |
| f"hidden={args.hidden}, relations=3) params={n_params:,}\n") |
|
|
| best_val_acc = 0.0 |
| checkpoint = Path(args.checkpoint) |
|
|
| print(f"{'Epoch':>5} {'Loss':>8} {'Val Acc':>8} {'':>6}") |
| print("-" * 35) |
|
|
| for epoch in range(1, args.epochs + 1): |
| loss = train_epoch(model, train_loader, optimizer, device, pos_weight) |
| val_acc = evaluate(model, valid_loader, device) |
| scheduler.step() |
|
|
| marker = "" |
| if val_acc > best_val_acc: |
| best_val_acc = val_acc |
| torch.save(model.state_dict(), checkpoint) |
| marker = "<- best" |
|
|
| print(f"{epoch:>5} {loss:>8.4f} {val_acc:>8.2%} {marker}") |
|
|
| print(f"\nLoading best checkpoint ({checkpoint}) ...") |
| model.load_state_dict(torch.load(checkpoint, map_location=device, |
| weights_only=True)) |
| test_acc = evaluate(model, test_loader, device) |
| print(f"Test accuracy: {test_acc:.2%}") |
| print(f"\nCheckpoint saved to: {checkpoint.resolve()}\n") |
|
|
| print("--- Results ---") |
| print(f" §7 Instruction-level GNN (full graph, 30ep h=64): 58.00%") |
| print(f" §4d Block-level GNN (full graph, 60ep h=128): 57.84%") |
| print(f" §11 Slice-GNN (backward DFG from sinks, this run): {test_acc:.2%}") |
|
|
| if test_acc >= 0.62: |
| print("\n Slice-GNN reaches VulPathFinder territory (61%) -- " |
| "signal concentration from slicing is the key improvement") |
| elif test_acc >= 0.58: |
| print("\n Improvement over instruction-level baseline -- " |
| "slicing reduces noise; consider R-GAT for further gains") |
| else: |
| print("\n No improvement over §7 -- slicing may be too aggressive " |
| "(too many functions have no detectable sinks) or too noisy " |
| "(fallback full graphs dominate). Check slice stats from preprocess_slice.py.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
