predict_SingleGUNet.py

from __future__ import print_function, division
import os
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
from equivariant_lib.models.G_UNet_model import G_UNet, G_RadioWNet
import random
import logging
import yaml
# Ignore warnings
import warnings
import torch.nn.functional as F

warnings.filterwarnings("ignore")
from lib import loaders, modules
import torch
import time
from collections import defaultdict
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 calc_loss(pred, target, metrics):
    batch_size = target.size(0)
    batch_mse = F.mse_loss(pred, target, reduction='mean')
    metrics['total_mse_sum'] += batch_mse.item() * batch_size

    batch_target_power = torch.mean(target ** 2)
    metrics['total_target_power_sum'] += batch_target_power.item() * batch_size

def print_metric(metrics, total_samples):
    # print_metrics_test(metrics, epoch_samples, 'mse')
    # Final RMSE: Sqrt of the average of all individual squared errors
    final_rmse = (metrics['total_mse_sum'] / total_samples) ** 0.5
    # Final NMSE: Total MSE divided by total target power
    final_nmse = metrics['total_mse_sum'] / (metrics['total_target_power_sum'] + 1e-8)
    final_mse = metrics['total_mse_sum'] / total_samples
    print(f"Final RMSE: {final_rmse:.4f}")
    print(f"Final NMSE: {final_nmse:.4f}")
    print(f"Final MSE: {final_mse:.5f}")
    print(f"Total data sample num: {total_samples}")

    logging.info(f"Final RMSE: {final_rmse:.4f}")
    logging.info(f"Final NMSE: {final_nmse:.4f}")
    logging.info(f"Final MSE: {final_mse:.5f}")
    logging.info(f"Total data sample num: {total_samples}")
    return final_rmse, final_nmse, final_mse

# 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 predict_model(model, dataloader):
    model.eval()  # Set model to evaluate mode
    # criterion = nn.MSELoss()

    metrics = {'total_mse_sum': 0.0, 'total_target_power_sum': 0.0}
    train_loss_list = []
    val_loss_list = []


    since = time.time()

    #metrics = defaultdict(float)
    epoch_samples = 0

    for batch_idx, (inputs, targets) in enumerate(dataloader):
        inputs = inputs.to(device)
        targets = targets.to(device)

        outputs = model(inputs)
        calc_loss(outputs, targets, metrics)
        epoch_samples += inputs.size(0)

    # epoch_loss = metrics['loss'] / epoch_samples
    # logging.info(f"total loss {epoch_loss}")
    print_metric(metrics, epoch_samples)

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

    return train_loss_list, val_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['predict']['log_filename']
    group = configs['first_gunet']['group']
    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['predict']['batch_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
    }

    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)
    model.load_state_dict(torch.load(os.path.join(model_path, group+'best_model.mdl'),
                                     map_location=torch.device(device)))
    pytorch_total_params = sum([np.prod(p.shape) for p in model.parameters()])
    logging.info(f"Total parameters of UNet: {pytorch_total_params}")

    logging.info('start to predict')
    predict_model(model, dataloaders['test'])