train_ResNet-IR.py

"""
Created on Mar 2023
@author: 
@project: EventSleep
"""

import torch, torchvision
from torchvision.models import ResNet18_Weights
from torch.utils.data import TensorDataset, DataLoader
from data_tools import *
from utils import *

from pathlib import Path
import numpy as np
import time

from sklearn.metrics import balanced_accuracy_score
from sacred import Experiment
import imageio.v3 as iio

import glob
import os
import datetime

def warn(*args, **kwargs):
    pass
import warnings
warnings.warn = warn

ex = Experiment('ClassificationDeepSleep')

@ex.config
def toy_data():
    toy_data = False

@ex.config
def create_folder(toy_data):
    train_configs = [1, 2, 3]

    if toy_data: root_dir = f'{Path(os.getcwd())}/Toy_Data/'
    else: root_dir = f'{Path(os.getcwd())}/DATA/'

    if train_configs == [1, 2, 3]:
        folder_name_configs = 'TrainAllConfigs'
    else:
        folder_name_configs = f'TrainConfig{train_configs[0]}'

    file_name = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    folder_path = os.path.join(f'{Path(os.getcwd())}/Models/ResNet-IR/{folder_name_configs}/{file_name}')
    if not os.path.exists(folder_path):  os.makedirs(folder_path)



@ex.config
def frames_post_processing(root_dir):
    inf_width = 250
    inf_height = 180
    frames_per_image = 3
    infrared_frame_folder = f'{root_dir}/Infrared/TRAIN'
    data_augmentation = False
    crop_bed = True


@ex.config
def labels_dicts():
    labels_dict = LabelsNames()
    labels_id = list(labels_dict.keys())
    labels_names = list(labels_dict.values())

@ex.config
def model_hyperparameters(frames_per_image):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    load_pretrained_weights = True
    save_checkpoints = True

    load_trained_weights = False
    checkpoint_path = ''

    in_channels = frames_per_image
    out_channels = 10

    n_epochs = 10
    lr = 1e-3
    batch_size = 32
    weight_decay = 0.01

    weights_labels = GiveWeightsToLabels()
    metric = torch.nn.CrossEntropyLoss(weight=weights_labels.to(device))


@ex.config
def training_test_strategy(toy_data):
    if toy_data:
        subjects_train = [5]
        subjects_val = [11]
    else:
        subjects_train = [1, 2, 3, 4, 5, 6, 7, 8, 10]
        subjects_val = [11]


@ex.capture
def extract_infrared_data(infrared_frame_folder, inf_width, inf_height, frames_per_image, subjects, configs,
                        batch_size, crop_bed, data_augmentation,  shuffle, device, dirs=None):

    X, y = None, None
    if dirs is None:
        dirs = []
        for subject in subjects:
            for c in configs:
                dirs.append(sorted(glob.glob(f'{infrared_frame_folder}/subject{subject:02}_config{c}/clip*_label*')))
        dirs = [item for sublist in dirs for item in sublist]

    for dir in dirs:
        X_dir = None
        for im_path in sorted(glob.glob(f'{dir}/*.png')):
            im = iio.imread(im_path)
            im = im[:, :, 0]
            im = im[np.newaxis, :, :, np.newaxis]
            if X_dir is None: X_dir = im
            else: X_dir = np.vstack([X_dir, im])

        X_dir = X_dir.astype('float32')

        if crop_bed:
            X_dir = CropBed(X_dir, 'Infrared', subject)

        if inf_width != X_dir.shape[2] or inf_height != X_dir.shape[1]:
            X_dir = ResizeInfraredFrames(X_dir, inf_height, inf_width)

        if X_dir.shape[0] - (frames_per_image - 1) < 0:
            continue
        X_dir = StackInfraredImages(X_dir, frames_per_image)

        if data_augmentation:
            X_dir = DataAugmentationInfrared(X_dir)

        y_dir = GetLabelFromDirName(dir, X_dir.shape[0], 'Infrared', data_augmentation)

        if X is None:
            X, y = X_dir, y_dir
        else:
            X, y = np.vstack([X, X_dir]), np.hstack([y, y_dir])

    X = X.astype('float32')
    X = X.transpose(0, 3, 1, 2)
    X, y = torch.from_numpy(X).to(device), torch.from_numpy(y).to(device)
    dataset = TensorDataset(X, y)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)

    return loader


@ex.capture
def load_model(in_channels, out_channels, lr, device, weight_decay, load_pretrained_weights):
    if load_pretrained_weights:
        resnet_model = torchvision.models.resnet18(weights=ResNet18_Weights.DEFAULT)
    else:
        resnet_model = torchvision.models.resnet18()

    model = MyResNet(resnet_model, in_channels, out_channels)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    return model.to(device), optimizer

@ex.capture
def train_model(model, optimizer, metric, subjects_train, subjects_val, n_epochs, save_checkpoints, batch_size, lr,
                weight_decay, train_configs, folder_path):

    t0 = time.time()
    val_loader = extract_infrared_data(subjects=subjects_val, configs=train_configs, shuffle=True)
    training_loader = extract_infrared_data(subjects=subjects_train, configs=train_configs,
                                            batch_size=batch_size, shuffle=True)
    acc_val_max = 0

    # Train
    t1 = time.time()
    for n_ep in range(n_epochs):
        n_batch = 0
        for inputs, targets in training_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = metric(outputs, targets)
            loss.backward()
            optimizer.step()
            n_batch += 1

        acc_train = balanced_accuracy_score(targets.cpu().detach().numpy(), outputs.argmax(-1).cpu().detach().numpy())
        outputs_val, targets_val = [], []
        for in_val, t_val in val_loader:
            out_val = model(in_val)
            outputs_val.append(out_val.cpu().detach())
            targets_val.append(t_val.cpu().detach())
        outputs_val, targets_val = torch.cat(outputs_val), torch.cat(targets_val)
        acc_val = balanced_accuracy_score(targets_val.numpy(), outputs_val.argmax(-1).numpy())
        print('[Epoch %d] Loss_train: %.3f, Acc_train: %.3f ||  Acc_val: %.3f' % (n_ep + 1, loss.item(), acc_train, acc_val))

        if save_checkpoints and acc_val > acc_val_max:
            checkpoint = {
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'lr': lr,
                'batch_size': batch_size,
                'weight_decay': weight_decay,
                'epoch': n_ep + 1,
                'loss': loss
            }
            acc_val_max = acc_val
            checkpoint_file = os.path.join(folder_path, f'checkpoint_accval{round(acc_val, 2)}.pth')
    torch.save(checkpoint, checkpoint_file)

    t2 = time.time()
    t_data = t1 - t0
    t_model = t2 - t1
    print("Model Training Finished")
    print("Time spent loading data:", t_data)
    print("Time spent training model:", t_model)
    return


@ex.automain
def run(folder_path):
    ex.commands["print_config"]()

    model, optimizer = load_model()
    train_model(model=model, optimizer=optimizer)
    ex.commands["save_config"](config_filename=f'{folder_path}/train_details.json')