JLB-JLB/android-ransomware-gnn-baseline

Android Ransomware GNN Baseline

Binary graph-level classifier for Android ransomware detection using Graph Neural Networks on function call graphs (FCGs).

Trained on the JLB-JLB/android-ransomware-fcg-baseline dataset (715 APKs: 502 benign, 213 ransomware).

Model Architecture

Parameter	Value
Architectures	GIN / GCN / GAT
GNN layers	3
Hidden dim	128
Pooling	Global mean pooling
Classifier	2-layer MLP with dropout=0.5
Parameters	~231K
Input features	5

Training Details

Parameter	Value
Optimizer	Adam (lr=0.001, wd=0.0001)
Scheduler	ReduceLROnPlateau (patience=10, factor=0.5)
Loss	Weighted cross-entropy (inverse class frequency)
Batch size	4
Max epochs	200
Early stopping	patience=20, min_epochs=100
Data split	70/15/15 stratified (seed=42)
Mixed precision	AMP (float16 on CUDA)

Results

Model	Dataset Variant	Accuracy	Macro F1	Malware F1	Precision	Recall	AUROC	Epochs	Time
GIN	internal_only	0.9815	0.9782	0.9697	0.9412	1.0000	0.9901	135	884s
GIN	full_fcg	0.9537	0.9468	0.9275	0.8649	1.0000	0.9774	113	935s
GCN	internal_only	0.9630	0.9556	0.9375	0.9375	0.9375	0.9836	108	1057s
GCN	full_fcg	0.9444	0.9345	0.9091	0.8824	0.9375	0.9885	101	1449s
GAT	internal_only	0.9630	0.9564	0.9394	0.9118	0.9688	0.9848	100	2047s
GAT	full_fcg	0.9815	0.9778	0.9688	0.9688	0.9688	0.9942	113	3646s

Repo Structure

README.md                        # This model card
model.py                         # Model class definitions with PyTorchModelHubMixin
GIN/
  internal_only/
    model.safetensors            # Trained weights
    config.json                  # Model __init__ kwargs (auto-saved by mixin)
    test_results.json            # Full test metrics
    training_log.json            # Per-epoch training history
  full_fcg/
    ...
GCN/
  internal_only/
    ...
  full_fcg/
    ...
GAT/
  internal_only/
    ...
  full_fcg/
    ...

Usage

Models use PyTorchModelHubMixin for native from_pretrained() support:

# Get the model class definition
from huggingface_hub import hf_hub_download
import importlib.util

mod_path = hf_hub_download("JLB-JLB/android-ransomware-gnn-baseline", "model.py")
spec = importlib.util.spec_from_file_location("model", mod_path)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)

# Download model weights to a local directory, then load
import tempfile
from pathlib import Path

with tempfile.TemporaryDirectory() as tmpdir:
    for fname in ["config.json", "model.safetensors"]:
        hf_hub_download(
            "JLB-JLB/android-ransomware-gnn-baseline", fname,
            subfolder="GIN/internal_only",
            local_dir=tmpdir,
        )
    model = model_module.GINClassifier.from_pretrained(
        str(Path(tmpdir) / "GIN" / "internal_only")
    )
model.eval()

# Download and load dataset
import torch
ds_path = hf_hub_download(
    "JLB-JLB/android-ransomware-fcg-baseline", "internal_only/fcg_dataset.pt", repo_type="dataset"
)
dataset = torch.load(ds_path, weights_only=False)

# Run inference
from torch_geometric.loader import DataLoader
batch = next(iter(DataLoader([dataset[0]], batch_size=1)))
with torch.no_grad():
    pred = model(batch.x, batch.edge_index, batch.batch).argmax(1).item()
    print(f"Prediction: {'benign' if pred == 0 else 'malware'}")

Requirements

torch>=2.0
torch-geometric>=2.4

Citation

If you use this model, please cite the associated thesis work.