train_SingleGUNet.py

from __future__ import print_function, division
import os
import pandas as pd
import yaml
from skimage import io, transform
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils, datasets, models
from equivariant_lib.models.G_UNet_model import G_UNet, G_RadioWNet
import random
import logging

# Ignore warnings
import warnings
warnings.filterwarnings("ignore")
from lib import loaders, modules

import torch
import torch.optim as optim
from torch.optim import lr_scheduler
import time
import copy
from collections import defaultdict
import torch.nn.functional as F
import torch.nn as nn


def calc_loss_dense(pred, target, metrics, criterion):
    
    loss = criterion(pred, target)
    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)

    return loss

# def calc_loss_sparse(pred, target, samples, metrics, num_samples):
#     criterion = nn.MSELoss()
#     loss = criterion(samples*pred, samples*target)*(256**2)/num_samples
#     metrics['loss'] += loss.data.cpu().numpy() * target.size(0)

#     return loss

def print_metrics(metrics, epoch_samples, phase):
    outputs1 = []
    outputs2 = []
    for k in metrics.keys():
        outputs1.append("{}: {:4f}".format(k, metrics[k] / epoch_samples))

    logging.info("{}: {}".format(phase, ", ".join(outputs1)))

def train_model(model, optimizer, scheduler, dataloaders, num_epochs=50, model_path='./data/', group=None):
    criterion = nn.MSELoss()
    # WNetPhase: traine first U and freez second ("firstU"), or vice verse ("secondU").
    # targetType: train against dense images ("dense") or sparse measurements ("sparse")
    #best_model_wts = copy.deepcopy(model.state_dict())
    best_loss = 1e10
    train_loss_list = []
    test_loss_list = []

    for epoch in range(num_epochs):
        logging.info('Epoch {}/{}'.format(epoch, num_epochs - 1))
        # print('-' * 10)

        since = time.time()

        # Each epoch has a training and validation phase
        for phase in ['train', 'test']:
            if phase == 'train':
                scheduler.step()
                # for param_group in optimizer.param_groups:
                #     print("learning rate", param_group['lr'])

                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            metrics = defaultdict(float)
            epoch_samples = 0

            for batch_idx,(inputs, targets) in enumerate(dataloaders[phase]):
                inputs = inputs.to(device)
                targets = targets.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                # with torch.set_grad_enabled(phase == 'train'):
                outputs = model(inputs)
                #outputs = normalize_output(outputs, inputs[:,0:1,:,:])
                loss = calc_loss_dense(outputs, targets, metrics, criterion)

                # backward + optimize only if in training phase
                if phase == 'train':
                    loss.backward()
                    optimizer.step()

                # statistics
                epoch_samples += inputs.size(0)
                # one_batch_loss = loss.item() / inputs.size(0)
                # logging.info(f"batch {batch_idx}, one batch loss: {one_batch_loss}")
            

            print_metrics(metrics, epoch_samples, phase)
            epoch_loss = metrics['loss'] / epoch_samples
            if phase == 'train':
                train_loss_list.append(epoch_loss)
            elif phase == 'test':
                test_loss_list.append(epoch_loss)

            # deep copy the model
            if phase == 'test' and epoch_loss < best_loss:
                best_loss = epoch_loss
                logging.info("saving best model")
                torch.save(model.state_dict(), os.path.join(model_path, group+'best_model.mdl'))
                

        time_elapsed = time.time() - since
        logging.info('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

    logging.info('Best test loss: {:4f}'.format(best_loss))

    # load best model weights
    # model.load_state_dict(torch.load(model_path, weights_only=True))
    # return model
    return train_loss_list, test_loss_list

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

if __name__ == '__main__':
    with open('configs/SingleGUNet.yaml') as f:
        configs = yaml.safe_load(f)
    log_dir = configs['log_dir']
    log_filename = configs['train']['log_filename']
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, log_filename)
    logging.basicConfig(
        filename=log_path,
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        filemode='w'  # 'a' for append (default), 'w' to overwrite each time
    )
    set_seed(516)
    batch_size = configs['train']['batch_size']
    num_epochs = configs['train']['num_epochs']
    lr = configs['train']['lr']
    schedule_step_size = configs['train']['schedule_step_size']
    model_path = configs['model_path']
    if not os.path.exists(model_path):
        os.makedirs(model_path)
    
    abs_working_dir = os.path.dirname(os.path.abspath(__file__))
    dataset_path = os.path.join(abs_working_dir,'RadioMapSeer') + '/'
    logging.info ('abs data path %s', dataset_path)
    
    Radio_train = loaders.RadioUNet_c(phase="train",dir_dataset=dataset_path)
    Radio_val = loaders.RadioUNet_c(phase="val",dir_dataset=dataset_path)
    Radio_test = loaders.RadioUNet_c(phase="test",dir_dataset=dataset_path)
    
    image_datasets = {
        'train': Radio_train, 'val': Radio_val, 'test': Radio_test
    }

    dataloaders = {
        'train': DataLoader(Radio_train, batch_size=batch_size, shuffle=False, num_workers=0),
        'val': DataLoader(Radio_val, batch_size=batch_size, shuffle=False, num_workers=0),
        'test': DataLoader(Radio_test, batch_size=batch_size, shuffle=False, num_workers=0)
    }

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    #os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
    torch.set_default_dtype(torch.float32)
    if device == 'cuda':
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        torch.backends.cudnn.enabled

    config_first_GUNet = configs['first_gunet']
    model = G_UNet (config_first_GUNet).to(device)
    pytorch_total_params = sum([np.prod(p.shape) for p in model.parameters()])
    logging.info(f"Total parameters of UNet: {pytorch_total_params}")


    optimizer_ft = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=schedule_step_size, gamma=0.1)
    logging.info ('start to train UNet')
    train_loss_first, test_loss_list = train_model(model, optimizer_ft, exp_lr_scheduler,dataloaders,
                                                   num_epochs=num_epochs,
                                                   model_path=model_path,
                                                   group = configs['first_gunet']['group'])
    first_loss_df = pd.DataFrame({'train_loss':train_loss_first, 'test_loss':test_loss_list})
    first_loss_df.to_csv(os.path.join(model_path,'first_unet_loss.csv'), index=False)