PyTTrain.py

import torch
from torch._C import device, dtype

from torch.functional import Tensor
from torch.serialization import save

from torch.utils.data import DataLoader
from torch import optim

from pathlib import Path

from torch.utils.tensorboard import SummaryWriter

from tqdm import tqdm

from PyTorchGTK.conf import SegIniReader, OptIniReader, IniReader

import optuna
import numpy as np
from datetime import datetime
from torch .utils import tensorboard as tb


class FitModel:

    def __init__(self, dL : DataLoader = None,
                        dLV : DataLoader = None,
                        model = None,
                        optF : optim.Optimizer = optim.SGD,
                        nEpochs : int = 5,
                        lr : float = 0.5,
                        startEp : int = 0,
                        storPath : Path = Path('.'),
                        logdir : Path = Path('logs'),
                        weightsName : str = 'weights',
                        dev = None,
                        dtype = torch.float32,
                        conf : IniReader | None = None) -> None:
        super().__init__()

        self.rs = np.random.RandomState(np.random.randint(1000))
        self.accuracy = {'Acc': lambda x , y: torch.Tensor(0.0)    }

        self.device = dev if not dev in [None, ''] else torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.dL = dL
        self.dLV = dLV
        self.model = model
        self.optF = optF
        self.nEpochs = nEpochs
        self.stEpoch = startEp
        self.ie = startEp
        self.lr = lr
        
        self.nCheck = 1
        self.bestAcc = {'Acc': -np.Inf}
        self.bestLoss = {'Acc': -np.Inf}

        self._ep_lossCalc = torch.zeros(1, device=self.device, requires_grad=False)
        self._ep_vAcc = torch.zeros(1, device=self.device, requires_grad=False)
        self._ep_vLos = torch.zeros(1, device=self.device, requires_grad=False)

        self._optimiser = None

        self.storPath = str(storPath)
        self.logdir = logdir
        self.wName = weightsName

        self.dtype = dtype
        print("training on ", self.device)

    def setLoaders(self, dL : DataLoader, dLV : DataLoader):
        self.dL = dL
        self.dLV = dLV

    def setModel(self, model, optF : optim.Optimizer, nEpochs : int, lr : float):
        self.model = model
        self.optF = optF
        self.nEpochs = nEpochs
        self.lr = lr

    def _prepTraining(self):

        self.tbWriter = SummaryWriter(self.logdir)

        self.bestAcc = {ik: -np.Inf for ik in self.accuracy.keys()}
        self.bestLoss = {ik: -np.Inf for ik in self.accuracy.keys()}

        self.lossFunc.to(self.device, dtype = self.dtype)
        self.model.to(self.device, dtype = self.dtype)
        # It must be after the model being sent to device.
        self._optimiser = self.optF(self.model.parameters(), lr = self.lr)
  
        """
        CHECK THAT IT WORKS !!
        """
        if not self.loadPath == None and Path(self.loadPath).exists():
            nend = self.nEpochs
            self.load(self.loadPath)
            # self._optimiser.param_groups[0]['lr'] = self.lr
            self.nEpochs = nend if self.nEpochs < nend else self.nEpochs + nend

    def prepBatch(self, b):
        return b

    def _ep_train(self):
        self.model.train()
        self._ep_lossCalc = torch.zeros(1, device=self.device, requires_grad=False)
        self._ep_acc = {ik : torch.zeros(1, device=self.device, requires_grad=False) for ik in self.accuracy.keys()}
        
        for b in tqdm(self.dL):
            xb, yb = self.prepBatch(b)
            xb = xb.to(self.device, dtype= self.dtype)
            yb = yb.to(self.device, dtype= self.dtype)

            self._optimiser.zero_grad()

            preds = self.model(xb)
            loss = self.lossFunc(preds, yb)

            self._ep_lossCalc += loss.detach()

            with torch.no_grad():
                for iacc in self.accuracy.keys():
                    self._ep_acc[iacc] += self.accuracy[iacc](preds, yb).detach()

            loss.backward()
            self._optimiser.step()
                # p.step()

        self._ep_lossCalc /= len(self.dL)
        for iacc in self.accuracy.keys():
            self._ep_acc[iacc] /= len(self.dL)
        return self._ep_lossCalc, self._ep_acc

    def _ep_Val(self):
        self._ep_vAcc = {ik : torch.zeros(1, device=self.device, requires_grad=False) for ik in self.accuracy.keys()}
        self._ep_vLos = torch.zeros(1, device=self.device, requires_grad=False)

        if len(self.dLV) == 0:
            return self._ep_vLos, self._ep_vAcc

        self.model.eval()
        with torch.no_grad():
            for b in tqdm(self.dLV):
                xv, yv = self.prepBatch(b)
                xv = xv.to(self.device)
                yv = yv.to(self.device)

                vPred = self.model(xv)
                loss = self.lossFunc(vPred, yv)
                self._ep_vLos += loss.detach()
                for iacc in self.accuracy.keys():
                    self._ep_vAcc[iacc] += self.accuracy[iacc](vPred, yv).detach()
        
        self._ep_vLos /= len(self.dLV)
        for iacc in self.accuracy.keys():
            self._ep_vAcc[iacc] /= len(self.dLV)
        return self._ep_vLos, self._ep_vAcc

    def _ep_end(self):
        with torch.set_grad_enabled(False):
            self._ep_lossCalc = self._ep_lossCalc.item()
            self._ep_acc = {ik : iv.item() for ik, iv in self._ep_acc.items()}
            self._ep_vLos = self._ep_vLos.item()
            self._ep_vAcc = {ik : iv.item() for ik, iv in self._ep_vAcc.items()}

            if self.ie % self.nCheck == 0:
                print("Epoch ", self.ie, "L-train: ", self._ep_lossCalc, " Acc-train: ", self._ep_acc,
                " L-val: ", self._ep_vLos, " Acc-val: ", self._ep_vAcc)

                self.tbWriter.add_scalars('Loss', {'train': self._ep_lossCalc, 'val': self._ep_vLos}, self.ie)
                for ik in self.accuracy:
                    self.tbWriter.add_scalars(ik, {'train': self._ep_acc[ik], 'val': self._ep_vAcc[ik]}, self.ie)

                    if self._ep_vAcc[ik] > self.bestAcc[ik]:
                        self.bestAcc[ik] = self._ep_vAcc[ik]
                        self.bestLoss[ik] = self._ep_lossCalc

                        self.save(Path(self.storPath) / (self.wName + '_' + ik + '.best'))


    def _end_training(self):
        with torch.set_grad_enabled(False):
            self.save(Path(self.storPath) / (self.wName + '.last'))
            self.tbWriter.flush()
            self.tbWriter.close()

    def lossFunc(self, x : Tensor, targ : Tensor) -> Tensor:
        return torch.sum(((x - targ)^2).view(-1))

    def fitModel(self, loadPath : str = None):
        self.loadPath = loadPath
        self._prepTraining()

        for self.ie in range(self.stEpoch, self.nEpochs):
            self._ep_train()
            self._ep_Val()
            self._ep_end()

        self._end_training()

    def load(self, fpath : Path) -> None:
        saved = torch.load(fpath, map_location=self.device)
        self.model.load_state_dict(saved['model_state_dict'])
        self.nEpochs = saved['nEpochs']
        self.stEpoch = saved['iEpoch'] + 1
        self.ie = self.stEpoch
        self._optimiser.load_state_dict(saved['opt_state_dict'])
        self.bestAcc = saved['Acc']
        self.bestLoss = saved['Loss']

    def save(self, fpath : Path) -> None:
        
        torch.save({
                'nEpochs': self.nEpochs,
                'iEpoch': self.ie,
                'opt_state_dict': self._optimiser.state_dict(),
                'model_state_dict': self.model.state_dict(),
                'Acc': self.bestAcc,
                'Loss': self.bestLoss
            }, fpath)

class FitModelMixedPrec(FitModel):
    def __init__(self, dL : DataLoader = None,
                        dLV : DataLoader = None,
                        model = None,
                        optF : optim.Optimizer = optim.SGD,
                        nEpochs : int = 5,
                        lr : float = 0.5,
                        startEp : int = 0,
                        storPath : Path = Path('.'),
                        logdir : Path = Path('logs'),
                        weightsName : str = 'weights',
                        dev = None,
                        dtype = torch.float32,
                        conf : IniReader | None = None) -> None:
        super().__init__(dL=dL, dLV=dLV, model=model, optF=optF, nEpochs=nEpochs, lr=lr, startEp=startEp, storPath=storPath, logdir=logdir, weightsName=weightsName, dev=dev, dtype=dtype)
        self.scaler = torch.cuda.amp.GradScaler()

    def _ep_train(self):
        self.model.train()
        self._ep_lossCalc = torch.zeros(1, device=self.device, requires_grad=False)
        self._ep_acc = {ik : torch.zeros(1, device=self.device, requires_grad=False) for ik in self.accuracy.keys()}

        for b in tqdm(self.dL):
            xb, yb = self.prepBatch(b)
            xb = xb.to(self.device, dtype= self.dtype)
            yb = yb.to(self.device, dtype= self.dtype)

            self._optimiser.zero_grad()
            with torch.cuda.amp.autocast():
                preds = self.model(xb)
                loss = self.lossFunc(preds, yb)

            self._ep_lossCalc += loss.detach()

            with torch.no_grad():
                for iacc in self.accuracy.keys():
                    self._ep_acc[iacc] += self.accuracy[iacc](preds, yb).detach()

            self.scaler.scale(loss).backward()
            self.scaler.step(self._optimiser)
            self.scaler.update()

        self._ep_lossCalc /= len(self.dL)
        for iacc in self.accuracy.keys():
            self._ep_acc[iacc] /= len(self.dL)
        return self._ep_lossCalc, self._ep_acc
   

from typing import Any

class OptimModel(object):

    def __init__(self, fitClass, conf : OptIniReader) -> None:
        self.conf = conf
        # Initialize the stepping process
        self.conf.__iter__()

        self.fitClass = fitClass

        self.dLoaderVal = None

        self.dataPath = Path(self.conf['DATASET']['root'])
        self.resFolder = self.dataPath / self.conf['DATASET']['store'] / ('res_' + self.conf['GENERAL']['name_id'])
        if not (self.dataPath / self.conf['DATASET']['store']).exists():
            (self.dataPath / self.conf['DATASET']['store']).mkdir()
        if not self.resFolder.exists():
            self.resFolder.mkdir()

        self.current_time = datetime.now().strftime('%b%d-%H-%M-%S')

        self.logsFold = self.resFolder / ('logs' + self.current_time)
        if not self.logsFold.exists(): self.logsFold.mkdir()

    def optimise(self, trial):
        self.trial = trial
        self.conf.trial = self.trial
        self.conf.__next__()

        icID = self.conf.getItId().replace('_', '-')

        dl, dlV = self.conf.prepObject('DATASET')

        device = self.conf['GENERAL']['device'] if not self.conf['GENERAL']['device'] in [None, ''] else torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        bSize = self.conf['DATASET']['batch']

        swloop = tb.writer.SummaryWriter(self.logsFold / icID)
        swloop.add_text('parameters/' + icID, self.conf.itStr().replace('\n', '<br />'), 0)

        self.acc2Check = iter(self.conf['TRAINING']['acc'].keys()).__next__()

        try:
            model = self.conf.getModel()
            model.to(device=device)
        except Exception:
            print('cannot make a model')
            raise optuna.exceptions.TrialPruned()

        loss, acc = self.conf.prepObject('TRAINING')

        opt = self.conf.getOptim()

        self.conf.writeConfig(self.resFolder)

        if 'learning_rate' in self.conf['TRAINING']:
            lrate = self.conf['TRAINING']['learning_rate']
        else:
            lrz = self.conf['TRAINING']['learning_rate_z']
            lrn = self.conf['TRAINING']['learning_rate_n']
            lrate = pow(10,-lrz) * lrn

        self.FM = self.fitClass(dl, dlV, model, optF = opt,
                nEpochs = self.conf['TRAINING']['epochs'], lr = lrate,
                storPath = self.resFolder / self.conf.getItId(), logdir = self.logsFold / (icID + '-plot'),
                dtype = self.conf['MODEL']['data_type'], conf = self.conf )
        self.FM.lossFunc = loss
        self.FM.accuracy = acc
        self._old_ep_end, self._old_prepTraining = self.FM._ep_end, self.FM._prepTraining
        self.FM._ep_end, self.FM._prepTraining = self._ep_end, self._prepTraining

        try:
            self.FM.fitModel()
        except Exception as ex:
            if hasattr(self.FM, 'bestAcc'):
                swloop.add_text('Result/' + icID, 'Acc: ' + str(self.FM.bestAcc) + ' Loss: ' + str(self.FM.bestLoss))
                with open(str(self.dataPath / self.conf['DATASET']['store'] / (self.conf['GENERAL']['name_id'] + '_optstd.txt')), 'a')  as fid:
                    fid.write(icID + ';' + self.current_time + ';' +
                                'Acc:' + str(self.FM.bestAcc) + ';Loss:' + str(self.FM.bestLoss) + ';' +
                                self.conf.itStr().replace('\n', ';') + 
                                '\n')
            swloop.flush()
            swloop.close()

            dl.dataset.__del__()
            dlV.dataset.__del__()
            dl, dlV = [], []
            model = []
            opt = []

            if type(ex) == optuna.exceptions.TrialPruned:
                raise ex
            if hasattr(self.FM, 'bestAcc'):
                return self.FM.bestAcc[self.acc2Check]
            elif hasattr(self.FM, '_ep_vAcc'):
                return self.FM._ep_vAcc[self.acc2Check]
            else:
                return 0.

        with open(str(self.dataPath / self.conf['DATASET']['store'] / (self.conf['GENERAL']['name_id'] + '_optstd.txt')), 'a')  as fid:
            fid.write(icID + ';' + self.current_time + ';' +
                        'Acc:' + str(self.FM.bestAcc) + ';Loss:' + str(self.FM.bestLoss) + ';' +
                        self.conf.itStr().replace('\n', ';') + 
                        '\n')
        swloop.add_text('Result/' + icID, 'Acc: ' + str(self.FM.bestAcc) + ' Loss: ' + str(self.FM.bestLoss))
        swloop.flush()
        swloop.close()

        dl.dataset.__del__()
        dlV.dataset.__del__()
        dl, dlV = [], []
        model = []
        opt = []

        return self.FM.bestAcc[self.acc2Check]

    def _ep_end(self):
        self._old_ep_end()

        try:
            self.trial.report(self.FM._ep_vAcc[self.acc2Check], self.FM.ie)
        except Exception:
            self.FM._end_training()
            raise optuna.exceptions.TrialPruned()
            
        if self.trial.should_prune():
            self.FM._end_training()
            raise optuna.exceptions.TrialPruned()

    def _prepTraining(self):
        torch.manual_seed(self.conf['GENERAL']['rand_seed'])
        self.rs = np.random.RandomState(self.conf['GENERAL']['rand_seed'])
        return self._old_prepTraining()