Pytorch_CNN-RNN/utils/train_methods.py

import torch

from torch import nn, optim
from sklearn.metrics import roc_curve, auc, confusion_matrix, classification_report
import seaborn as sns


# GENERAL PURPOSE
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import time


# TRAIN
def train(model, train_data, test_data, CNN_filepath, epochs=20, graphs=True):
    model.train()
    criterion = nn.CrossEntropyLoss(reduction='mean')
    optimizer = optim.Adam(model.parameters(), lr=1e-5)

    losses = pd.DataFrame(columns=['Epoch', 'Avg_loss', 'Time'])
    start_time = time.time()  # seconds

    for epoch in range(epochs):  # loop over the dataset multiple times
        epoch += 1

        # Estimate & count training time
        t = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
        t_remain = time.strftime("%H:%M:%S", time.gmtime((time.time() - start_time)/epoch * epochs))
        print(f"{epoch/epochs * 100} || {epoch}/{epochs} || Time: {t}/{t_remain}")

        running_loss = 0.0

        # Batches & training
        for i, data in enumerate(train_data, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data[0].to(model.device), data[1].to(model.device)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = model.forward(inputs)
            loss = criterion(outputs, labels)   # This loss is the mean of losses for the batch
            loss.backward()
            optimizer.step()

            # adds average batch loss to running loss
            running_loss += loss.item()

            # mini-batches for progress
            if(i%10==0 and i!=0):
                print(f"{i}/{len(train_data)}, temp. loss:{running_loss / len(train_data)}")

        # average loss
        avg_loss = running_loss / len(train_data)      # Running_loss / number of batches
        print(f"Avg. loss: {avg_loss}")

        # loss on validation
        val_loss = evaluate(test_data, graphs)

        losses = losses.append({'Epoch':int(epoch), 'Avg_loss':avg_loss, 'Val_loss':val_loss, 'Time':time.time() - start_time}, ignore_index=True)


    print('Finished Training')
    losses.to_csv('./cnn_net_data.csv')

    if(graphs):
        # MAKES EPOCH VS AVG LOSS GRAPH
        plt.plot(losses['Epoch'], losses['Avg_loss'], label="Loss on Training")
        plt.xlabel('Epoch')
        plt.ylabel('Average Loss')
        plt.title('Loss vs Epoch On Training & Validation data')

        # MAKES EPOCH VS VALIDATION LOSS GRAPH
        plt.plot(losses['Epoch'], losses['Val_loss'], label="Loss on Validation")
        plt.savefig('./avgloss_epoch_curve.png')
        plt.show()

        torch.save(model.state_dict(), CNN_filepath)
        print("Model saved")



def load(model, filepath):
    model.load_state_dict(torch.load(filepath))


def evaluate(model, val_data, graphs=True):
# EVALUATE MODEL
    correct = 0
    total = 0

    predictionsLabels = []
    predictionsProbabilities = []
    true_labels = []

    criterion = nn.CrossEntropyLoss(reduction='mean')
    model.eval()
    # since we're not training, we don't need to calculate the gradients for our outputs
    with torch.no_grad():
        for data in val_data:
            images, labels = data[0].to(model.device), data[1].to(model.device)
            # calculate outputs by running images through the network
            outputs = model.forward(images)
            # the class with the highest energy is what we choose as prediction

            loss = criterion(outputs, labels)  # mean loss from batch

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

            # Saves predictionsProbabilities and labels for ROC
            if(graphs):
                predictionsLabels.extend(predicted.cpu().numpy())
                predictionsProbabilities.extend(outputs.data[:, 1].cpu().numpy())     # Grabs probability of positive
                true_labels.extend(labels.cpu().numpy())

    print(f'Accuracy of the network on {total} scans: {100 * correct // total}%')

    if(not graphs): print(f'Validation loss: {loss.item()}')
    else:
        # ROC
        # Calculate TPR and FPR
        fpr, tpr, thresholds = roc_curve(true_labels, predictionsProbabilities)

        # Calculate AUC
        roc_auc = auc(fpr, tpr)

        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC: {roc_auc})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.005])
        plt.ylim([0.0, 1.005])

        plt.xlabel('False Positive Rate (1 - Specificity)')
        plt.ylabel('True Positive Rate (Sensitivity)')
        plt.title('Receiver Operating Characteristic (ROC) Curve')
        plt.legend(loc="lower right")
        plt.savefig('./ROC.png')
        plt.show()


        # Calculate confusion matrix
        cm = confusion_matrix(true_labels, predictionsLabels)

        # Plot confusion matrix
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
        plt.xlabel('Predicted labels')
        plt.ylabel('True labels')
        plt.title('Confusion Matrix')
        plt.savefig('./confusion_matrix.png')
        plt.show()

        # Classification Report
        report = classification_report(true_labels, predictionsLabels)
        print(report)

    model.train()

    return(loss.item())


# PREDICT
def predict(model, data):
    model.eval()
    with torch.no_grad():
        for data in data:
            images, labels = data[0].to(model.device), data[1].to(model.device)
            outputs = model.forward(images)
            # the class with the highest energy is what we choose as prediction
            _, predicted = torch.max(outputs.data, 1)
    model.train()
    return (labels, predicted)  # RETURNS (true, predicted)