You can deploy this model and start earning money today!

XRailgunX/
malware-cnn-malimg

Name: malware-cnn-malimg
Availability: InStock

Views

No views yet

Malware Classification CNN on Malimg

Trained models for 25-class malware family classification on the Malimg dataset. Four checkpoints from a phase-based optimization study, going from a baseline CNN (89.06%) to an EfficientNetB0-based model (98.48% with TTA).

GitHub (code, reports, training curves): github.com/ffftuanxxx/malware-classification-CNN-optimized

Checkpoints

File	Phase	Architecture	Val Accuracy	Macro F1	Size
`baseline/pesi.h5`	Baseline	3-block Conv + Flatten + Dense(256) (59M params)	89.06%	86.45%	227 MB
`phase1/best_model.h5`	Phase 1	Conv + BN + GAP + Dense (<1M params)	67.28%	39.28%	5.2 MB
`phase2/best_model.h5`	Phase 2	EfficientNetB0 + Dense head, two-stage fine-tune	94.91%	83.95%	30 MB
`phase3/best_model.h5`	Phase 3	Phase 2 + Focal Loss + oversampling + Cosine + TTA	98.48%	95.78%	30 MB

All metrics are on the 923-sample Malimg val split.

How to Load

python

from huggingface_hub import hf_hub_download
import tensorflow as tf

# Download any checkpoint
path = hf_hub_download(repo_id="XRailgunX/malware-cnn-malimg",
                       filename="phase3/best_model.h5")

# Phase 1/2/3 saved full models; load directly
model = tf.keras.models.load_model(path, compile=False)

# Baseline is weights-only; rebuild architecture first
# (see run_malimg_classifier.py in the GitHub repo), then:
# model.load_weights(path)

For Phase 3, which uses a custom Focal Loss, pass compile=False when loading (the loss function is not serialized). Then recompile if you plan to continue training.

Input Format

Input shape: (256, 256, 3) RGB (grayscale Malimg images replicated to 3 channels for ImageNet-pretrained backbones).
Preprocessing:
- Baseline / Phase 1: rescale to [0, 1] (x / 255.0).
- Phase 2 / Phase 3 (EfficientNet-based): keep raw [0, 255] float — EfficientNet has a built-in Normalization layer. Applying rescale=1./255 externally will cause predictions to collapse.

Classes

25 Malimg malware families (class indices in alphabetical order): Adialer.C, Agent.FYI, Allaple.A, Allaple.L, Alueron.gen!J, Autorun.K, C2LOP.P, C2LOP.gen!g, Dialplatform.B, Dontovo.A, Fakerean, Instantaccess, Lolyda.AA1, Lolyda.AA2, Lolyda.AA3, Lolyda.AT, Malex.gen!J, Obfuscator.AD, Rbot!gen, Skintrim.N, Swizzor.gen!E, Swizzor.gen!I, VB.AT, Wintrim.BX, Yuner.A.

Training Details (Phase 3, best model)

Backbone: EfficientNetB0 (ImageNet pretrained)
Head: GlobalAveragePooling2D → Dense(256, relu) + BN + Dropout(0.5) → Dense(25, softmax)
Stage 1 (10 epochs, frozen base, Adam 1e-3): 68% → 91% val acc
Stage 2 (30 epochs, top 20 layers unfrozen, Cosine decay from 1e-4 to 1.3e-6): best val_loss 0.01090 at epoch 24
Loss: Focal Loss (γ=2.0, α=0.25)
Augmentation: width/height_shift=0.1, horizontal_flip=True, zoom=0.1, brightness=[0.9,1.1] (NO vertical flip, rotation, or color jitter)
Minority oversampling: bootstrap each class with <200 samples to 200
TTA at inference: 5 random augmentations averaged → +1.41pp over single-view prediction

Training curves, per-class F1, and full technical report are in the GitHub repo.

Per-class F1 (Phase 3)

Perfect (F1 = 1.00) on 19 of 25 classes. Remaining:

Class	Support	F1
Autorun.K	10	0.86
C2LOP.P	14	0.86
C2LOP.gen!g	20	0.95
Swizzor.gen!E	12	0.77
Swizzor.gen!I	13	0.53
Yuner.A	80	0.98

Swizzor.gen!I is the hardest remaining case — a near-visual duplicate of Swizzor.gen!E. Single-model upper bound on Malimg seems to be around this level; reaching 99%+ typically requires multi-model ensembles or byte-level auxiliary features.

Limitations

Trained and evaluated only on Malimg (25 PE malware families, grayscale byte images). Does not transfer to mobile malware, scripts, or obfuscated packers outside this distribution.
Not calibrated for out-of-distribution detection. Softmax confidence on non-Malimg inputs is unreliable.
Evaluation is on the val split that was also used for model selection, so the reported numbers are slightly optimistic vs a held-out test run.

Citation

If you build on this work, please also cite the upstream baseline repo cridin1/malware-classification-CNN and the Malimg dataset (Nataraj et al., 2011).