> For the complete documentation index, see [llms.txt](https://aashraymt.gitbook.io/docs/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://aashraymt.gitbook.io/docs/sms-spam-detector/malware-classification.md).

# Malware Classification

This all started as a weekend ML project challenge I came across.

When diving into malware analysis, one of the most frustrating bottlenecks is the sheer amount of time manual classification takes. Figuring out whether a sample belongs to a specific threat family like Emotet or WannaCry usually demands an exhaustive mix of static and dynamic analysis. You end up deep in the weeds, reverse engineering binaries just to pull out behavioral traits, delivery methods, and technical signatures.

I wanted to tackle this challenge by exploring a way to significantly speed up this pipeline while keeping the workflow secure. That led me to the core concept of this project: applying Machine Learning to automate malware classification.

Specifically, I decided to build a classifier based on malware images. Turning malicious binaries into images to train a Convolutional Neural Network (CNN) might sound a bit counterintuitive at first, but the advantages are massive. By converting the raw binary data into a visual format, I can train the model on the structural patterns of different malware families without ever having to handle or detonate the live executables directly. It completely eliminates the risk of accidentally infecting the host system, making it a perfect, sandbox safe approach for a lab environment while still delivering highly accurate threat categorization.

## 1. Theory: Malware Visualization & CNNs <a href="#id-1-theory-malware-visualization-cnns" id="id-1-theory-malware-visualization-cnns"></a>

This project classifies malware families from images instead of directly handling executable binaries. The approach is based on the concepts explored in the paper *Malware Classification with Deep Learning*.

#### The Key Idea <a href="#the-key-idea" id="the-key-idea"></a>

A PE (Portable Executable) binary is, at its core, a sequence of bytes. You can:

* Read these bytes as integers ranging from 0 to 255.
* Arrange them into a 2D array.
* Treat that array as a grayscale image.

Samples from the same threat family tend to share similar byte level structures. When rendered visually, these structural similarities show up as consistent textures and patterns. Convolutional Neural Networks (CNNs) are incredibly efficient at picking up local patterns and textures in images, making them the perfect tool to learn and distinguish between malware families based on these visual signatures.

#### How Binary to Image Conversion Works <a href="#how-binary-to-image-conversion-works" id="how-binary-to-image-conversion-works"></a>

The conversion process is entirely straightforward:

1. Read the malware binary as raw bytes.
2. Map each byte (0 to 255) into an 8 bit vector.
3. Arrange the bytes into a 2D grid.
4. Render the grid as a grayscale image.

Every single binary byte is fully encoded within the image. This means the image can be used to exactly reconstruct the malicious binary without any loss of information. As long as you treat the file as raw bytes, map each byte directly to one pixel, and avoid lossy operations like compression or color changes, you can take any malware, convert it to an image, and convert it right back securely.

<figure><img src="/files/DHB32CWzBYUNoEYm4iPU" alt=""><figcaption></figcaption></figure>

**Note:** In our actual training pipeline, we resize the images to 75x75 pixels to fit the CNN architecture, which means the model itself works on a lossy representation. The lossless property applies strictly to the original visualization, not the training input.

## 2. Getting and Exploring the Dataset <a href="#id-2-getting-and-exploring-the-dataset" id="id-2-getting-and-exploring-the-dataset"></a>

#### Downloading and Unpacking <a href="#downloading-and-unpacking" id="downloading-and-unpacking"></a>

We are using the `malimg` dataset, which contains 9,339 image files distributed across 25 different malware families. Each image is a direct visual representation of a Windows PE file.

```
wget https://www.kaggle.com/api/v1/datasets/download/ikrambenabd/malimg-original -O malimg.zip
unzip malimg.zip
```

#### Exploring the Dataset <a href="#exploring-the-dataset" id="exploring-the-dataset"></a>

Before building anything, we need to understand our data. We set up the base path and count the number of samples per malware family:

```
import os
import matplotlib.pyplot as plt
import seaborn as sns

DATA_BASE_PATH = "./malimg_paper_dataset_imgs/"

# Compute the class distribution
dist = {}
for mlw_class in os.listdir(DATA_BASE_PATH):
    mlw_dir = os.path.join(DATA_BASE_PATH, mlw_class)  # Constructs the full path
    dist[mlw_class] = len(os.listdir(mlw_dir))
```

#### Visualizing Class Distribution <a href="#visualizing-class-distribution" id="visualizing-class-distribution"></a>

To spot any potential biases in our data, we graph it out using custom styling:

```
# dataclasses = list(dist.keys())
frequencies = list(dist.values())

# plot
plt.figure(facecolor=node_black)
sns.barplot(y=classes, x=frequencies, edgecolor="black", orient='h', color=htb_green)
plt.title("Malware Class Distribution", color=htb_green)
plt.xlabel("Malware Class Frequency", color=htb_green)
plt.ylabel("Malware Class", color=htb_green)
plt.show()
```

<figure><img src="/files/UenwcEeOf9O5f7Jo9Zol" alt=""><figcaption></figcaption></figure>

When you look at the resulting graph, you will notice some malware families appear far more frequently than others. This class imbalance could potentially skew the model and produce inaccurate predictions. In a production environment, you would want to fine tune the dataset before training to ensure a more balanced class distribution.

## 3. Preprocessing and Splitting <a href="#id-3-preprocessing-and-splitting" id="id-3-preprocessing-and-splitting"></a>

#### Splitting with split-folders <a href="#splitting-with-split-folders" id="splitting-with-split-folders"></a>

We need distinct datasets for training and testing to prove our model actually learns rather than just memorizes. We use the `split-folders` library to handle this automatically:

```
pip3 install split-folders
```

```
import splitfolders

DATA_BASE_PATH = "./malimg_paper_dataset_imgs/"
TARGET_BASE_PATH = "./newdata/"
TRAINING_RATIO = 0.8
TEST_RATIO = 1 - TRAINING_RATIO

splitfolders.ratio(input=DATA_BASE_PATH, output=TARGET_BASE_PATH, ratio=(TRAINING_RATIO, 0, TEST_RATIO))
```

This creates a `newdata/` directory with `train/` and `test/` subdirectories. Each contains the exact same 25 family subfolders, properly separated via an 80/20 split.

#### Installing PyTorch <a href="#installing-pytorch" id="installing-pytorch"></a>

Since we do not need the massive overhead of the full CUDA 12 stack for this phase, we install the official CPU wheel index. This keeps the environment lightweight and saves a ton of disk space:

```
pip3 install --index-url https://download.pytorch.org/whl/cpu torch==2.2.2 torchvision==0.17.2
```

## 4. Normalization and DataLoaders <a href="#id-4-normalization-and-dataloaders" id="id-4-normalization-and-dataloaders"></a>

#### Why Normalize? <a href="#why-normalize" id="why-normalize"></a>

Normalization ensures our images are standardized. Essentially, we are taking an image and turning it into a mathematical tensor so the neural network can actually process it.

Here is what the transformation does:

* Resizes the image to 75x75 pixels.
* Converts the image matrix into a PyTorch Tensor.
* Standardizes the pixel values using established ImageNet mean and standard deviation metrics.

```
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((75, 75)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
```

#### Applying the Transform to All Images <a href="#applying-the-transform-to-all-images" id="applying-the-transform-to-all-images"></a>

The following code creates PyTorch dataset objects. These objects know exactly how to crawl your directory structure, apply the transformation rules, and return the `(image_tensor, label)` pairs whenever you index them:

```
from torchvision.datasets import ImageFolder
import os

BASE_PATH = "./newdata/"

train_dataset = ImageFolder(
    root=os.path.join(BASE_PATH, "train"),
    transform=transform
)

test_dataset = ImageFolder(
    root=os.path.join(BASE_PATH, "test"),
    transform=transform
)
```

#### Creating DataLoaders <a href="#creating-dataloaders" id="creating-dataloaders"></a>

A PyTorch `DataLoader` is an incredible helper tool. Instead of manually loading files one by one, the DataLoader creates an easy to use iterator over batches of data.

```
from torch.utils.data import DataLoader

train_loader = DataLoader(
    train_dataset,
    batch_size=TRAIN_BATCH_SIZE,
    shuffle=True,
    num_workers=2
)

test_loader = DataLoader(
    test_dataset,
    batch_size=TEST_BATCH_SIZE,
    shuffle=False,
    num_workers=2
)
```

Here is exactly what the DataLoader handles for you under the hood:

* **Batching:** Instead of processing 1 image at a time, you process batches (like 1024 images at once). This is significantly faster and aligns with how neural networks execute matrix math.
* **Shuffling:** Setting `shuffle=True` randomizes the order during every epoch so the model does not overfit to a fixed sequence.
* **Parallel Loading:** Setting `num_workers=2` means two background processes are actively loading and preprocessing images while your model trains on the current batch. This eliminates CPU bottlenecking.
* **Clean Training Loops:** Your actual training code can just execute `for batch in train_loader:` without worrying about file paths or batching logic.

#### Checking the Normalized Image <a href="#checking-the-normalized-image" id="checking-the-normalized-image"></a>

If you want to verify the pipeline, you can pull a normalized sample directly from the DataLoader:

```
sample = next(iter(train_loader))[0][0]
plt.imshow(sample.permute(1, 2, 0))
plt.show()
```

<figure><img src="/files/DVFkH7hZIAiwKFkATXRI" alt=""><figcaption></figcaption></figure>

#### Combining Everything into a Single Function <a href="#combining-everything-into-a-single-function" id="combining-everything-into-a-single-function"></a>

For clean code architecture, we wrap all of this into a single reusable loader function:

```
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import os

def load_datasets(base_path, train_batch_size, test_batch_size):
    transform = transforms.Compose([
        transforms.Resize((75, 75)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    train_dataset = ImageFolder(
        root=os.path.join(base_path, "train"),
        transform=transform
    )

    test_dataset = ImageFolder(
        root=os.path.join(base_path, "test"),
        transform=transform
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=train_batch_size,
        shuffle=True,
        num_workers=2
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=test_batch_size,
        shuffle=False,
        num_workers=2
    )

    n_classes = len(train_dataset.classes)
    return train_loader, test_loader, n_classes
```

## 5. The Model: Transfer Learning with ResNet50 <a href="#id-5-the-model-transfer-learning-with-resnet50" id="id-5-the-model-transfer-learning-with-resnet50"></a>

Building a network from scratch takes massive amounts of time and computing power. Instead, we use transfer learning with a pretrained ResNet50 model. We download the pretrained weights, run our dataset through it, and fine tune the architecture for our specific security use case.

#### Class Definition and Loading ResNet50 <a href="#class-definition-and-loading-resnet50" id="class-definition-and-loading-resnet50"></a>

```
import torch.nn as nn
import torchvision.models as models

HIDDEN_LAYER_SIZE = 1000

class MalwareClassifier(nn.Module):
    def __init__(self, n_classes):
        super(MalwareClassifier, self).__init__()
        # Load pretrained ResNet50
        self.resnet = models.resnet50(weights='DEFAULT')
```

Here, `MalwareClassifier` inherits from PyTorch's `nn.Module`. When we initialize it, `models.resnet50(weights='DEFAULT')` pulls down a 50 layer deep network already trained on ImageNet. It starts with a massive baseline understanding of spatial textures and visual patterns.

#### Freezing All ResNet Layers <a href="#freezing-all-resnet-layers" id="freezing-all-resnet-layers"></a>

```
                for param in self.resnet.parameters():
            param.requires_grad = False
```

A model learns by updating its weights via gradients. By setting `requires_grad = False`, we explicitly tell PyTorch to freeze the entire backbone. The result:

* Gradients are not computed for these layers.
* The optimizer ignores them completely.
* Training becomes incredibly fast because we are only training the final layer, bypassing the need to update 23 million parameters.

#### Replacing the Last Fully Connected Layer <a href="#replacing-the-last-fully-connected-layer" id="replacing-the-last-fully-connected-layer"></a>

```
        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Linear(num_features, HIDDEN_LAYER_SIZE),
            nn.ReLU(),
            nn.Linear(HIDDEN_LAYER_SIZE, n_classes)
        )
```

The `self.resnet.fc` variable represents the final classification layer. We grab the incoming feature count (2048 for ResNet50) and replace the generic output with a custom sequential head mapping 2048 features down to our 25 specific malware classes.

#### Forward Pass <a href="#forward-pass" id="forward-pass"></a>

```
    def forward(self, x):
        return self.resnet(x)
```

This defines the execution pathway. When you pass a batch of images to the model, it flows through the frozen convolutional layers to extract features, and then hits our brand new final layer to generate the class predictions.

#### Initializing the Model <a href="#initializing-the-model" id="initializing-the-model"></a>

```
train_loader, test_loader, n_classes = load_datasets(DATA_PATH, TRAINING_BATCH_SIZE, TEST_BATCH_SIZE)
model = MalwareClassifier(n_classes)
```

## 6. Training <a href="#id-6-training" id="id-6-training"></a>

An epoch in machine learning signifies one complete pass of the entire training dataset through a model. This pass allows the model to calculate its errors, update its internal weights, and learn the underlying patterns.

#### Setup: Loss, Optimizer and Bookkeeping <a href="#setup-loss-optimizer-and-bookkeeping" id="setup-loss-optimizer-and-bookkeeping"></a>

```
import torch
import time

def train(model, train_loader, n_epochs, verbose=False):
    model.train()
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters())

    training_data = {"accuracy": [], "loss": []}
```

* `model.train()` activates training features like dropout and batch normalization.
* `CrossEntropyLoss()` is the standard loss function for multi class scenarios. It compares the model's outputs against the true labels and calculates a mathematical penalty.
* `Adam()` is the optimizer that actively tweaks the unfrozen weights to minimize that loss penalty.

#### Outer Loop: Repeat for Each Epoch <a href="#outer-loop-repeat-for-each-epoch" id="outer-loop-repeat-for-each-epoch"></a>

```
        for epoch in range(n_epochs):
        running_loss = 0
        n_total = 0
        n_correct = 0
        checkpoint = time.time() * 1000
```

For every epoch, we initialize trackers for the running loss, the total samples processed, the number of correct predictions, and a timestamp to benchmark execution speed.

#### Inner Loop: Go Over All Batches <a href="#inner-loop-go-over-all-batches" id="inner-loop-go-over-all-batches"></a>

```
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
```

This is the absolute core of the machine learning process:

1. **Iterate:** Grab a batch of images and their corresponding labels.
2. **Clear:** `optimizer.zero_grad()` wipes the gradients from the previous batch to prevent stacking errors.
3. **Forward Pass:** `outputs = model(inputs)` generates the predictions.
4. **Calculate Error:** `loss = criterion()` figures out how wrong the predictions were.
5. **Backpropagation:** `loss.backward()` calculates the gradients across all trainable parameters.
6. **Optimize:** `optimizer.step()` applies those gradients to adjust the network weights.

#### Measuring Accuracy and Accumulating Stats <a href="#measuring-accuracy-and-accumulating-stats" id="measuring-accuracy-and-accumulating-stats"></a>

```
            _, predicted = outputs.max(1)
            n_total += labels.size(0)
            n_correct += predicted.eq(labels).sum().item()
            running_loss += loss.item()
```

We pull the index of the highest probability prediction using `outputs.max(1)`. We then compare it against the true label, tally up the correct guesses, and add the batch loss to our running total.

#### End of Epoch: Compute Averages, Log and Print <a href="#end-of-epoch-compute-averages-log-and-print" id="end-of-epoch-compute-averages-log-and-print"></a>

```
        epoch_loss = running_loss / len(train_loader)
        epoch_duration = int(time.time() * 1000 - checkpoint)
        epoch_accuracy = compute_accuracy(n_correct, n_total)

        training_data["accuracy"].append(epoch_accuracy)
        training_data["loss"].append(epoch_loss)

        if verbose:
            print(f"[i] Epoch {epoch+1} of {n_epochs}: Acc: {epoch_accuracy:.2f}% "
                  f"Loss: {epoch_loss:.4f} (Took {epoch_duration} ms).")

    return training_data
```

#### Training Results <a href="#training-results" id="training-results"></a>

```
[i] Epoch 1 of 10: Acc: 60.60% Loss: 1.4193 (Took 28579 ms).
[i] Epoch 2 of 10: Acc: 87.30% Loss: 0.4063 (Took 25845 ms).
[i] Epoch 3 of 10: Acc: 90.79% Loss: 0.2605 (Took 26132 ms).
[i] Epoch 4 of 10: Acc: 92.00% Loss: 0.2161 (Took 26910 ms).
[i] Epoch 5 of 10: Acc: 93.79% Loss: 0.1864 (Took 26168 ms).
[i] Epoch 6 of 10: Acc: 94.74% Loss: 0.1497 (Took 26220 ms).
[i] Epoch 7 of 10: Acc: 95.20% Loss: 0.1344 (Took 25920 ms).
[i] Epoch 8 of 10: Acc: 95.47% Loss: 0.1294 (Took 25816 ms).
[i] Epoch 9 of 10: Acc: 96.18% Loss: 0.1112 (Took 25490 ms).
[i] Epoch 10 of 10: Acc: 96.74% Loss: 0.0997 (Took 25784 ms).
[i] Inference accuracy: 88.83%.
```

<figure><img src="/files/Dohr1ZT47e1UMTM9nG9b" alt=""><figcaption></figcaption></figure>

<figure><img src="/files/19AZNUtGuoG1UipxtKEX" alt=""><figcaption></figcaption></figure>

As the epochs progress, the training loss steadily drops while the accuracy climbs. The model successfully maps visual anomalies directly to the associated threat classifications.

## 7. Evaluation and Inference <a href="#id-7-evaluation-and-inference" id="id-7-evaluation-and-inference"></a>

#### Predict Function <a href="#predict-function" id="predict-function"></a>

This function isolates a single piece of input and returns the network's top prediction:

```
def predict(model, test_data):
    model.eval()

    with torch.no_grad():
        output = model(test_data)
        _, predicted = torch.max(output.data, 1)

    return predicted
```

Switching to `model.eval()` ensures inference behaves consistently, and `torch.no_grad()` shuts off the gradient engine entirely since we are just testing, not training.

#### Evaluate Function <a href="#evaluate-function" id="evaluate-function"></a>

We loop this prediction mechanism over the entire, unseen test dataset to get our final real world accuracy metrics:

```
def compute_accuracy(n_correct, n_total):
    return round(100 * n_correct / n_total, 2)

def evaluate(model, test_loader):
    model.eval()

    n_correct = 0
    n_total = 0

    with torch.no_grad():
        for data, target in test_loader:
            predicted = predict(model, data)
            n_total += target.size(0)
            n_correct += (predicted == target).sum().item()

    accuracy = compute_accuracy(n_correct, n_total)

    return accuracy
```

The loop steps through every batch in `test_loader`, generates the predictions, and checks them against the target labels. Ultimately, the model proves its reliability with an **88.83% test accuracy**.