model.py

import os, time
import pandas as pd

import torch
from torch import nn

import networks
from data import recovery_df


def mkdirs(path):
    if not os.path.exists(path):
        os.makedirs(path)


class RMEP(object):
    r"""the RMEP model"""
    def __init__(self, opt):
        #super(RMEP, self).__init__()
        self.opt = opt
        self.isTrain = not opt.isTest
        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')  # get device name: CPU or GPU
        
        # save all the checkpoints to save_dir
        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
        mkdirs(self.save_dir)

        self.dataroot = opt.dataroot

        # criterion and output nc
        if self.opt.model_type == 'reg':
            self.output_nc = len(opt.output_params)
            self.criterion = nn.SmoothL1Loss()
            self.output_params = opt.output_params
        elif (self.opt.model_type == 'clf') or (self.opt.model_type == 'clf_multi'):
            self.criterion = nn.CrossEntropyLoss()
            self.output_nc = opt.output_nc
            self.output_params = ['label']
        else:
            raise ValueError(f'Unknown model type {opt.model_type}.')

        # get the neural network
        self.netRMEP = networks.define_net(
            opt.model_name, opt.input_nc, self.output_nc, opt.nrf, opt.norm, opt.init_type, opt.init_gain, opt.num_blk)
        self.netRMEP.to(self.device)
        self.model_names = ['RMEP']

        # get the optimizer        
        self.optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.netRMEP.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))
            
    def set_input(self, input):
        r"""Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): includes the data itself and its metadata information.
        """
        self.X = input[0].to(self.device, non_blocking=True)
        self.y = input[1].to(self.device, non_blocking=True)
        self.name = input[2]
    
    def setup(self):
        r"""Load and print networks; create schedulers

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        self.scheduler = networks.get_scheduler(self.optimizer, self.opt)
        if not self.isTrain or self.opt.continue_train:
            self.load_networks(self.opt.epoch)
        self.print_networks(self.opt.verbose)
    
    def forward(self):
        self.out = self.netRMEP(self.X)
        #print(self.out, self.out.shape,  self.y, self.y.shape)

    def optimize_parameters(self):
        self.optimizer.zero_grad()
        self.loss = self.criterion(self.out, self.y)
        self.loss.backward()
        self.optimizer.step()    

    def train_one_epoch(self, loader):
        tic = time.time()
        data_times = 0
        loss = 0
        tic_data = time.time()
        self.netRMEP.train()
        for data in loader:
            data_times += time.time() - tic_data
            self.set_input(data)
            self.forward()
            self.optimize_parameters()
            loss += self.loss.item()
            tic_data = time.time()
        return loss/len(loader), time.time()-tic, data_times/len(loader)

    @torch.no_grad()
    def test_one_epoch(self, loader, savename=None):
        tic = time.time()
        data_times = 0
        loss = 0
        tic_data = time.time()
        self.netRMEP.eval()

        metric = []
        for idx, data in enumerate(loader):
            data_times += time.time() - tic_data
            self.set_input(data)
            self.forward()
            loss += self.criterion(self.out, self.y).item()
            if not self.isTrain:
                if self.opt.model_type == 'clf' or self.opt.model_type == 'clf_multi':
                    _, pred = torch.max(self.out, 1)
                elif self.opt.model_type == 'reg':
                    pred = self.out
                metric.append(self.name + self.y.flatten().cpu().numpy().tolist() + pred.flatten().cpu().numpy().tolist())            
            tic_data = time.time()
        
        if not self.isTrain:            
            columns =  ['name'] + [item for item in self.output_params] + ['pred_'+item for item in self.output_params]
            df = pd.DataFrame(metric, columns=columns)            
            if self.opt.model_type == 'reg':
                df = recovery_df(df, self.output_params)
            df.to_csv(os.path.join(self.save_dir, f'{savename}.csv'), index=False)
        return loss/len(loader), time.time()-tic, data_times/len(loader)

    def update_learning_rate(self):
        r"""Update learning rates for all the networks; called at the end of every epoch"""
        old_lr = self.optimizer.param_groups[0]['lr']
        #for scheduler in self.schedulers:
        if self.opt.lr_policy == 'plateau':
            self.scheduler.step(0.0)
        else:
            self.scheduler.step()
        lr = self.optimizer.param_groups[0]['lr']
        #print('learning rate %.7f -> %.7f' % (old_lr, lr))
        return old_lr, lr

    def save_networks(self, epoch):
        r"""Save all the networks to the disk.
        Parameters:
            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
        """
        for name in self.model_names:
            if isinstance(name, str):
                save_filename = 'net%s_%s.pth' % (name, epoch)
                save_path = os.path.join(self.save_dir, save_filename)
                net = getattr(self, 'net' + name)

                if torch.cuda.is_available():
                    torch.save(net.cpu().state_dict(), save_path)
                    net.cuda()
                else:
                    torch.save(net.cpu().state_dict(), save_path)

    def load_networks(self, epoch):
        r"""Load all the networks from the disk.
        Parameters:
            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
        """
        for name in self.model_names:
            if isinstance(name, str):
                load_filename = 'net%s_%s.pth' % (name, epoch)
                load_path = os.path.join(
                    self.save_dir, load_filename)
                net = getattr(self, 'net' + name)
                print('loading the model from %s' % load_path)
                
                state_dict = torch.load(
                    load_path, map_location=str(self.device))
                if hasattr(state_dict, '_metadata'):
                    del state_dict._metadata

                net.load_state_dict(state_dict)

    def print_networks(self, verbose):
        r"""Print the total number of parameters in the network and (if verbose) network architecture
        Parameters:
            verbose (bool) -- if verbose: print the network architecture
        """
        print('---------- Networks initialized -------------')
        for name in self.model_names:
            if isinstance(name, str):
                net = getattr(self, 'net' + name)
                num_params = 0
                for param in net.parameters():
                    num_params += param.numel()
                if verbose:
                    print(net)
                print('[Network %s] Total number of parameters : %.3f M' %
                      (name, num_params / 1e6))
        print('-----------------------------------------------')


def build_model(opt):
    r"""
    build the network
    """
    return RMEP(opt)