train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import models as models
from dataset_preparation import MalwareDetect2Dataset, OpcodeDataset
from tqdm import tqdm
import yaml
import numpy as np
from typing import Optional


class ModelParams:
    def __init__(
        self,
        input_size: int = 16 * 16,
        hidden_size: Optional[int] = 256,
        hidden_size1: Optional[int] = 128,
        hidden_size2: Optional[int] = 64,
        input_channels: int = 1,
        num_classes: int = 5,
        num_epochs: int = 32,
        batch_size: int = 24,
        learning_rate: float = 0.001,
        patience: Optional[int] = None,
    ):
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.hidden_size1 = hidden_size1
        self.hidden_size2 = hidden_size2
        self.input_channels = input_channels
        self.num_classes = num_classes
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.patience = patience

class EarlyStopping:
    def __init__(self, patience=5, delta=0.0):
        """
        Early stopping to stop the training when the validation loss doesn't improve.

        Parameters:
        - patience (int): How many epochs to wait before stopping when no improvement.
        - delta (float): Minimum change in the monitored metric to qualify as an improvement.
        """
        self.patience = patience
        self.delta = delta
        self.best_loss = float('inf')
        self.counter = 0
        self.early_stop = False

    def __call__(self, val_loss):
        if val_loss < self.best_loss - self.delta:
            self.best_loss = val_loss
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True

# Function to train the model and report accuracy for each epoch
def train_model(model, criterion, optimizer, train_loader, num_epochs=10, device='cpu'):
    """
    Train the model and report loss and accuracy for each epoch.

    Parameters:
    - model (nn.Module): The neural network to train.
    - criterion (torch.nn.Module): Loss function.
    - optimizer (torch.optim.Optimizer): Optimizer for model parameters.
    - train_loader (DataLoader): DataLoader for training data.
    - device (torch.device): Device to run the model on (CPU or GPU).
    - num_epochs (int): Number of epochs to train.

    Returns:
    - dict: Training history containing 'loss' and 'accuracy' per epoch.
    """
    
    model.to(device)
    model.train()
    history = {'loss': [], 'accuracy': []}

    for epoch in range(num_epochs):
        running_loss = 0.0
        correct = 0
        total = 0

        epoch_start = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs}", unit="batch")

        for images, labels in epoch_start:
            images, labels = images.to(device), labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Accumulate training loss.
            running_loss += loss.item()

            # Calculate accuracy.
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        # Compute metrics for the epoch.
        epoch_loss = running_loss / len(train_loader)
        epoch_accuracy = 100 * correct / total

        history['loss'].append(epoch_loss)
        history['accuracy'].append(epoch_accuracy)

        # Log epoch results.
        print(f"Epoch [{epoch + 1}/{num_epochs}] - Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%")

    return history

def train_with_validation(model, criterion, optimizer, train_loader, val_loader, num_epochs=10, device='cpu'):
    """
    Train the model and report loss and accuracy for each epoch.

    Parameters:
    - model (nn.Module): The neural network to train.
    - criterion (torch.nn.Module): Loss function.
    - optimizer (torch.optim.Optimizer): Optimizer for model parameters.
    - train_loader (DataLoader): DataLoader for training data.
    - val_loader (DataLoader): DataLoader for validation data.
    - device (torch.device): Device to run the model on (CPU or GPU).
    - num_epochs (int): Number of epochs to train.

    Returns:
    - dict: Training history containing 'loss' and 'accuracy' per epoch.
    """
    
    model.to(device)
    model.train()
    history = {'loss': [], 'accuracy': [], 'val_loss': [], 'val_accuracy': []}

    for epoch in range(num_epochs):
        running_loss = 0.0
        correct = 0
        total = 0

        epoch_start = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs}", unit="batch")

        for images, labels in epoch_start:
            images, labels = images.to(device), labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Accumulate training loss.
            running_loss += loss.item()

            # Calculate accuracy.
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        # Compute metrics for the epoch.
        epoch_loss = running_loss / len(train_loader)
        epoch_accuracy = 100 * correct / total

        # Log epoch results.
        history['loss'].append(epoch_loss)
        history['accuracy'].append(epoch_accuracy)

        # Validate the model.
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item()

                _, predicted = torch.max(outputs, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()

        # Compute metrics for the epoch.
        epoch_val_loss = val_loss / len(val_loader)
        epoch_val_accuracy = 100 * val_correct / val_total

        history['val_loss'].append(epoch_val_loss)
        history['val_accuracy'].append(epoch_val_accuracy)

        print(f"Epoch [{epoch + 1}/{num_epochs}] - Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, Val-loss: {epoch_val_loss:.4f}, Val-Accuracy: {epoch_val_accuracy:.2f}%")

def train_with_validation_and_early_stopping(model, criterion, optimizer, train_loader, val_loader, num_epochs=10, device='cpu', patience=5):
    """
    Train the model with early stopping and report loss and accuracy for each epoch.

    Parameters:
    - model (nn.Module): The neural network to train.
    - criterion (torch.nn.Module): Loss function.
    - optimizer (torch.optim.Optimizer): Optimizer for model parameters.
    - train_loader (DataLoader): DataLoader for training data.
    - val_loader (DataLoader): DataLoader for validation data.
    - device (torch.device): Device to run the model on (CPU or GPU).
    - num_epochs (int): Number of epochs to train.
    - patience (int): Number of epochs to wait for validation loss improvement.

    Returns:
    - dict: Training history containing 'loss' and 'accuracy' per epoch.
    """

    model.to(device)
    history = {'loss': [], 'accuracy': [], 'val_loss': [], 'val_accuracy': []}

    early_stopping = EarlyStopping(patience=patience)

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        epoch_start = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs}", unit="batch")

        for images, labels in epoch_start:
            images, labels = images.to(device), labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward pass and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Accumulate training loss
            running_loss += loss.item()

            # Calculate accuracy
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        # Compute metrics for the epoch
        epoch_loss = running_loss / len(train_loader)
        epoch_accuracy = 100 * correct / total

        # Log epoch results
        history['loss'].append(epoch_loss)
        history['accuracy'].append(epoch_accuracy)

        # Validate the model
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item()

                _, predicted = torch.max(outputs, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()

        # Compute metrics for the epoch
        epoch_val_loss = val_loss / len(val_loader)
        epoch_val_accuracy = 100 * val_correct / val_total

        history['val_loss'].append(epoch_val_loss)
        history['val_accuracy'].append(epoch_val_accuracy)

        print(f"Epoch [{epoch + 1}/{num_epochs}] - Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, "
              f"Val-loss: {epoch_val_loss:.4f}, Val-Accuracy: {epoch_val_accuracy:.2f}%")

        # Check for early stopping
        early_stopping(epoch_val_loss)
        if early_stopping.early_stop:
            print(f"Early stopping triggered after epoch {epoch + 1}")
            break

    return history

def train_model(model, model_name, parameters, optimizer, train_loader, valid_loader, device):
    print("Hello World!")

def save_model_and_params(model, model_name, parameters, optimizer, history):
    print("Hello World!")

# Parameters
input_size = 16 * 16
hidden_size = 256
hidden_size1 = 128
hidden_size2 = 64
input_channels = 1
num_classes = 5
num_epochs = 32
batch_size = 24
learning_rate = 0.001
patience = 3

with open("config.yaml") as f:
    config = yaml.safe_load(f)

data = config["data"]["opcode_data_tfidf"]
opcode_frequency_data_zipped = np.load(data)
X_train = opcode_frequency_data_zipped['X_train']
y_train = opcode_frequency_data_zipped['y_train']

train_dataset = OpcodeDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# Create train and val datasets.
# train_dir = "./data/malware_detect2_simhash_1616_images/train"
# train_dataset = MalwareDetect2Dataset(data_dir=train_dir, transform=transform_cnn)
# train_size = int(0.9 * len(train_dataset))
# val_size = len(train_dataset) - train_size
# train_dataset, val_dataset = random_split(train_dataset, [train_size, val_size])

# Create train and val data loaders.
# train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
# val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)

# Initialize the model, criterion, and optimizer.
# model = models.CNN3DOFC1LayerNet().to(device)
model = models.LinearClassifier(input_size=393, num_classes=5).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

history = train_model(model, criterion, optimizer, train_loader, num_epochs, device)
# history = train_with_validation(model, criterion, optimizer, train_loader, val_loader, num_epochs, device)
# history = train_with_validation_and_early_stopping(model, criterion, optimizer, train_loader, val_loader, num_epochs, device, patience)


# Save the model's weights, biases, and hyperparameters.
save_data = {
    'model_state_dict': model.state_dict(),
    'input_size': input_size,
    # 'hidden_size': hidden_size,
    # 'hidden_size1': hidden_size1,
    # 'hidden_size2': hidden_size2,
    # 'input_channels': input_channels,
    'num_classes': num_classes,
    'num_epochs': num_epochs,
    'batch_size': batch_size,
    'learning_rate': learning_rate,
    # 'patience': patience,
    'history': history
}
torch.save(save_data, './models/linear_classifier_tfidf.pth')

# Optional: Plot training history
# plt.figure(figsize=(10, 5))
# plt.plot(history['loss'], label='Loss')
# plt.plot(history['accuracy'], label='Accuracy')
# plt.xlabel('Epoch')
# plt.ylabel('Metric')
# plt.legend()
# plt.title('Training Metrics')
# plt.show()

def main():
    params = ModelParams(393, None, None, None, None, 5, 32, 24, 0.001, None)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    with open("config.yaml") as f:
        config = yaml.safe_load(f)
        data = config["data"]["opcode_data"]
        opcode_frequency_data_zipped = np.load(data)
        X_train = opcode_frequency_data_zipped['X_train']
        y_train = opcode_frequency_data_zipped['y_train']

    train_dataset = OpcodeDataset(X_train, y_train)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

    model = models.LinearClassifier().to(device)

    criterion = nn.CrossEntropyLoss()

    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    history = train_model(model, criterion, optimizer, train_loader, num_epochs, device)


if __name__ == "__main__":
    main()