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

Training Details

Results

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.