train_Temp.py

import torch
import torch.optim as optim
from tqdm import tqdm
from models import RFACNN, RFAUNet, RFAAttUNet
from data_loader_Temp import TemperatureDataset, load_data, DataLoader, batch_size
from utils import new_combined_loss
from config import num_epochs, batch_size, alpha, beta, gamma, model_path_Temp, file_paths, model_name
from torch.optim.lr_scheduler import ReduceLROnPlateau
import matplotlib.pyplot as plt
import os

def get_model(choice):
    if choice == "1":
        return RFACNN()
    elif choice == "2":
        return RFAUNet()
    elif choice == "3":
        return RFAAttUNet()
    else:
        raise ValueError(f"Unknown model choice: {choice}")

class EarlyStopping:
    def __init__(self, patience=5, verbose=False, delta=0):
        self.patience = patience
        self.verbose = verbose
        self.delta = delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False
        self.best_model = None

    def __call__(self, val_loss, model):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.best_model = model.state_dict()
        elif val_loss > self.best_loss - self.delta:
            self.counter += 1
            if self.verbose:
                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.best_model = model.state_dict()
            self.counter = 0

def train_model(model, criterion, optimizer, train_loader, valid_loader, num_epochs):
    model.train()
    
    train_loss = []
    valid_loss = []

    early_stopping = EarlyStopping(patience=10, verbose=True)
    
    # Train loop
    for epoch in range(num_epochs):
        epoch_loss = 0

        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}", leave=False):
            Ninput_train_data_batch, MR_train_data_batch, Temp_train_data_batch = batch

            # Move data to GPU if available
            if torch.cuda.is_available():
                Ninput_train_data_batch = Ninput_train_data_batch.unsqueeze(1).cuda()
                MR_train_data_batch     = MR_train_data_batch.unsqueeze(1).cuda()
                Temp_train_data_batch   = Temp_train_data_batch.unsqueeze(1).cuda()

            # Forward pass
            outputs = model(Ninput_train_data_batch, MR_train_data_batch)

            # Calculate loss
            loss = criterion(outputs, Temp_train_data_batch, alpha, beta, gamma)

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

            # Accumulate the loss
            epoch_loss += loss.item()

        # Print the average loss for the epoch
        print(f'Epoch: {epoch+1}, Average Loss: {epoch_loss / len(train_loader)}')
        train_loss.append(epoch_loss / len(train_loader))
        
        # Validation loop
        model.eval()  # Set the model to evaluation mode
        with torch.no_grad():  # No gradients required for validation
            val_loss = 0

            for val_batch in valid_loader:
                # Forward pass
                Ninput_val_data_batch, MR_val_data_batch, Temp_val_data_batch = val_batch
                Temp_val_data_batch   = Temp_val_data_batch.unsqueeze(1).cuda()
                MR_val_data_batch     = MR_val_data_batch.unsqueeze(1).cuda()
                Ninput_val_data_batch = Ninput_val_data_batch.unsqueeze(1).cuda()

                val_outputs = model(Ninput_val_data_batch, MR_val_data_batch)

                # Calculate the loss
                loss = criterion(val_outputs, Temp_val_data_batch,alpha,beta,gamma)
                val_loss += loss.item()

            # Calculate average validation loss and update the scheduler
            avg_val_loss = val_loss / len(valid_loader)
            scheduler.step(avg_val_loss)

            # Print validation loss
            print(f'Epoch: {epoch+1}, Validation Loss: {val_loss / len(valid_loader)}')    
            valid_loss.append(val_loss / len(valid_loader))

            # Early stopping
            early_stopping(avg_val_loss, model)
            
            if early_stopping.early_stop:
                print("Early stopping")
                break

    # Load the best model
    model.load_state_dict(early_stopping.best_model)
    
    return train_loss, valid_loss
    

if __name__ == "__main__":
    # Load data from data_loader
    Temp_train_data, Ninput_train_data, MR_train_data, Temp_valid_data, Ninput_valid_data, MR_valid_data, Temp_test_data_foreseen, Temp_test_data_unforeseen, Ninput_test_data_foreseen, Ninput_test_data_unforeseen, MR_test_data_foreseen, MR_test_data_unforeseen = load_data(file_paths)

    # Initialize the model
    model = get_model(model_name)
    model.cuda() if torch.cuda.is_available() else model.cpu()

    # Define the loss function and optimizer
    criterion = new_combined_loss
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=5, verbose=True)

    # Load the training dataset
    Temp_train_dataset = TemperatureDataset(Ninput_train_data, MR_train_data, Temp_train_data)    
    train_loader = DataLoader(Temp_train_dataset, batch_size=batch_size, shuffle=False)
    
    # Load the validation dataset
    Temp_valid_dataset = TemperatureDataset(Ninput_valid_data, MR_valid_data, Temp_valid_data)    
    valid_loader = DataLoader(Temp_valid_dataset, batch_size=batch_size, shuffle=False)    

    # Train the model
    train_loss, valid_loss = train_model(model, criterion, optimizer, train_loader, valid_loader, num_epochs)
   
    # Save the trained model
    model_name = model.__class__.__name__
    torch.save(model.state_dict(), f'{model_path_Temp}/{model_name}_Temp_{num_epochs}epoch.pth')

    # Save train_loss and valid_loss to text files
    train_loss_file = f'train_graph/{model_name}_Temp_train_loss.txt'
    valid_loss_file = f'train_graph/{model_name}_Temp_valid_loss.txt'

    # Write train_loss to a text file
    with open(train_loss_file, 'w') as f:
        for loss in train_loss:
            f.write(f"{loss}\n")
        print(f"Training loss saved to {train_loss_file}")

    # Write valid_loss to a text file
    with open(valid_loss_file, 'w') as f:
        for loss in valid_loss:
            f.write(f"{loss}\n")
        print(f"Validation loss saved to {valid_loss_file}")

    # Set font sizes
    plt.rcParams.update({'font.size': 20}) 
    #plt.rcParams['axes.titlesize'] = 18
    #plt.rcParams['axes.labelsize'] = 16
    #plt.rcParams['xtick.labelsize'] = 14
    #plt.rcParams['ytick.labelsize'] = 14
    #plt.rcParams['legend.fontsize'] = 14

    # Create a training graph folder
    directory = "train_graph"
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"Directory '{directory}' created.")
    else:
        print(f"Directory '{directory}' already exists.")
        
    # Plot training graph   
    plt.figure(figsize=(10,5))
    plt.title("Training and Validation Loss")
    plt.plot(valid_loss,label="val")
    plt.plot(train_loss,label="train")
    plt.xlabel("iterations")
    plt.ylabel("Loss")
    plt.legend()
    plt.savefig(f'train_graph/{model_name}_Temp_train_valid.png', dpi=600, pad_inches=0)
    plt.close()