TrixyZ/trainer/trainer_scripts/pytorch_trainer.py

import torch
import torch.nn as nn
import torch.optim as optim
import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from utils.training_interface import TrainerInterface
from utils.file_utils import list_files_in_directory
import time
import json

# Global variables for progress tracking
current_progress = 0
current_loss = 0.0
current_accuracy = 0.0
total_epochs = 0
total_batches = 0

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

class PyTorchTrainer(TrainerInterface):
    def __init__(self, learning_rate=0.001):
        self.model_path = None
        self.training_directories = []
        self.sample_directory = None
        self.background_directory = None
        self.dropout_rate = 0.5
        self.epochs = 3000
        self.batch_size = 32
        self.learning_rate = learning_rate
        self.model = None
        self.optimizer = None
        self.criterion = nn.CrossEntropyLoss()
        self.current_epoch = 0
        self.stop_training = False
        self.training_status = "Not started"
        self.train_loader = None
        self.val_loader = None

    def set_model_path(self, model_path):
        self.model_path = model_path
        self.model = self.load_model()
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)

    def set_training_directories(self, directories):
        self.training_directories = directories

    def set_sample_directory(self, directory):
        self.sample_directory = directory

    def set_background_directory(self, directory):
        self.background_directory = directory

    def set_dropout_rate(self, rate):
        self.dropout_rate = rate

    def set_epochs(self, epochs):
        self.epochs = epochs

    def set_batch_size(self, batch_size):
        self.batch_size = batch_size

    def get_training_data(self):
        return self.training_directories

    def get_model_path(self):
        return self.model_path

    def get_epochs(self):
        return self.epochs

    def get_batch_size(self):
        return self.batch_size

    def get_dropout_rate(self):
        return self.dropout_rate

    def get_training_directories(self):
        return self.training_directories

    def get_sample_directory(self):
        return self.sample_directory

    def get_background_directory(self):
        return self.background_directory

    def set_training_data(self, features, labels):
        X_train, X_val, y_train, y_val = train_test_split(features, labels, test_size=0.2, random_state=42)
        self.train_loader = torch.utils.data.DataLoader(CustomDataset(X_train, y_train), batch_size=self.batch_size, shuffle=True)
        self.val_loader = torch.utils.data.DataLoader(CustomDataset(X_val, y_val), batch_size=self.batch_size, shuffle=False)

    def preprocess_training_data(self):
        import librosa
        data_path_dict = {}
        data_path_dict[0] = list_files_in_directory(self.background_directory, extension='.wav')
        for idx, directory in enumerate(self.training_directories):
            data_path_dict[idx+1] = list_files_in_directory(directory, extension='.wav')

        all_data = []
        total_files = sum(len(files) for files in data_path_dict.values())
        processed_files = 0

        for class_label, list_of_files in data_path_dict.items():
            for single_file in list_of_files:
                try:
                    audio, sample_rate = librosa.load(single_file)
                    mfcc = librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40)
                    mfcc_processed = np.mean(mfcc.T, axis=0)
                    all_data.append([mfcc_processed, class_label])
                    processed_files += 1
                    global current_progress
                    current_progress = (processed_files / total_files) * 100
                except Exception as e:
                    print(f"Exception: {str(e)}")

        df = pd.DataFrame(all_data, columns=["feature", "class_label"])
        df.to_pickle(self.model_path + "audio_data.csv")

    def load_model(self):
        num_classes = len(self.training_directories) + 1  # Assuming one extra for background or unknown
        model = nn.Sequential(
            nn.Linear(40, 128),
            nn.ReLU(),
            nn.Dropout(self.dropout_rate),
            nn.Linear(128, num_classes),  # Update output layer to match num_classes
            nn.Softmax(dim=1)
        )
        model_path = self.model_path + "hotword_model.pth"
        if os.path.exists(model_path):
            state_dict = torch.load(model_path)
            try:
                model.load_state_dict(state_dict)
            except RuntimeError as e:
                print(f"Error loading model state dict: {e}")
                # Handle dimension mismatch by reinitializing the last layer
                model = nn.Sequential(
                    nn.Linear(40, 128),
                    nn.ReLU(),
                    nn.Dropout(self.dropout_rate),
                    nn.Linear(128, num_classes),
                    nn.Softmax(dim=1)
                )
                print("Reinitialized the final layer to match the number of classes.")
        return model

    def start(self):
        self.training_status = "Training"
        self.stop_training = False
        self.run_training()

    def start_new(self):
        self.training_status = "Training"
        self.stop_training = False
        self.model = self.load_model()
        self.current_epoch = 0
        self.run_training()

    def pause(self):
        self.training_status = "Paused"
        self.stop_training = True

    def resume(self):
        self.training_status = "Training"
        self.stop_training = False
        self.run_training()

    def stop(self):
        self.training_status = "Stopped"
        self.stop_training = True

    def run_training(self):
        global current_progress, current_loss, current_accuracy, total_epochs, total_batches
        self.preprocess_training_data()
        start_time = time.time()
        df = pd.read_pickle(self.model_path + "audio_data.csv")
        X = df["feature"].values
        X = np.concatenate(X, axis=0).reshape(len(X), 40)
        y = np.array(df["class_label"].tolist())
        self.set_training_data(X, y)

        total_batches = len(self.train_loader)
        total_epochs = self.epochs
        
        training_info = {
            "trainer_script": "tensorflow_trainer",
            "training_time": 0,
            "model_size": 0,
            "epochs": self.epochs,
            "dropout_rate": self.dropout_rate,
            "batch_size": self.batch_size,
            "total_training_data": len(X),
            "training_directories": self.training_directories,
            "accuracy": [],
            "loss": [],
            "test_accuracy": 0
        }

        for epoch in range(self.current_epoch, self.epochs):
            if self.stop_training:
                break
            self.current_epoch = epoch
            self.model.train()
            for batch_idx, (batch_X, batch_y) in enumerate(self.train_loader):
                global current_loss, current_accuracy
                self.optimizer.zero_grad()
                outputs = self.model(batch_X)
                loss = self.criterion(outputs, batch_y.long())  # Convert labels to Long type
                loss.backward()
                self.optimizer.step()
                # Update global progress variables
                current_loss = loss.item()
                current_accuracy = (outputs.argmax(dim=1) == batch_y).float().mean().item()
                current_progress = (epoch * len(self.train_loader) + batch_idx + 1) / (self.epochs * len(self.train_loader)) * 100
            training_info["accuracy"].append(current_accuracy)
            training_info["loss"].append(current_loss)
            
            print(f"Epoch {epoch + 1}/{self.epochs}, Loss: {current_loss}, Accuracy: {current_accuracy}")

        epoch_time = time.time() - start_time
        training_info["training_time"] = epoch_time
        
        if not self.stop_training:
            torch.save(self.model.state_dict(), self.model_path + "hotword_model.pth")
            self.evaluate_model(X, y)
            self.save_training_info(training_info)
            self.training_status = "Finished"
        else:
            self.training_status = "Paused"

    def evaluate_model(self, X_test, y_test):
        self.model.eval()
        with torch.no_grad():
            X_test_tensor = torch.FloatTensor(X_test)
            y_test_tensor = torch.LongTensor(y_test)
            outputs = self.model(X_test_tensor)
            _, predicted = torch.max(outputs, 1)
            accuracy = (predicted == y_test_tensor).sum().item() / len(y_test_tensor)
            print(f"Test Accuracy: {accuracy * 100:.2f}%")
    
    def save_training_info(self, info):
        with open(self.model_path + "info_pytorch_trainer.json", "w") as f:
            json.dump(info, f)

    def get_training_info(self):
        return {
            "status": self.training_status,
            "current_epoch": self.current_epoch,
            "total_epochs": self.epochs,
            "loss": current_loss,
            "accuracy": current_accuracy
        }