trainer.py

from lightning_model import Wavegrad2
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from omegaconf import OmegaConf as OC
import os
import argparse
import datetime
from glob import glob
from pytorch_lightning.callbacks.base import Callback
import torch
from pytorch_lightning.utilities import rank_zero_only
from copy import deepcopy
from utils.tblogger import TensorBoardLoggerExpanded

# Other DDPM/Score-based model applied EMA
# In our works, there are no significant difference
class EMACallback(Callback):
    def __init__(self, filepath, alpha=0.999, k=3):
        super().__init__()
        self.alpha = alpha
        self.filepath = filepath
        self.k = 3 #max_save
        self.queue = []
        self.last_parameters = None

    @rank_zero_only
    def _del_model(self, removek):
        if os.path.exists(self.filepath.format(epoch=removek)):
            os.remove(self.filepath.format(epoch=removek))

    @rank_zero_only
    def on_train_batch_start(self, trainer, pl_module,batch, batch_idx,dataloader_idx):
        if hasattr(self, 'current_parameters'):
            self.last_parameters = self.current_parameters
        else:
            self.last_parameters = deepcopy(pl_module.state_dict())

    @rank_zero_only
    def on_train_batch_end(self, trainer, pl_module,outputs, batch, batch_idx,dataloader_idx):
        self.current_parameters = deepcopy(pl_module.state_dict())
        for k, v in self.current_parameters.items():
            self.current_parameters[k].copy_(self.alpha * v +
                                             (1. - self.alpha) *
                                             self.last_parameters[k])
        del self.last_parameters
        return

    @rank_zero_only
    def on_epoch_end(self, trainer, pl_module):
        if hasattr(self, 'current_parameters'):
            self.queue.append(trainer.current_epoch)
            torch.save(self.current_parameters,
                       self.filepath.format(epoch=trainer.current_epoch))
            pl_module.print(
                f'{self.filepath.format(epoch = trainer.current_epoch)} is saved')

            while len(self.queue) > self.k:
                self._del_model(self.queue.pop(0))

        else:
            self.current_parameters = deepcopy(pl_module.state_dict())

        return


def train(args):
    hparams = OC.load('hparameter.yaml')
    now = datetime.datetime.now().strftime('%m_%d_%H')
    hparams.name = f"{hparams.log.name}_{now}"
    os.makedirs(hparams.log.tensorboard_dir, exist_ok=True)
    os.makedirs(hparams.log.checkpoint_dir, exist_ok=True)
    model = Wavegrad2(hparams)
    tblogger = TensorBoardLoggerExpanded(hparams)
    ckpt_path = f'{hparams.log.name}_{now}_{{epoch}}'
    checkpoint_callback = ModelCheckpoint(dirpath=hparams.log.checkpoint_dir,
                                          filename=ckpt_path,
                                          verbose=True,
                                          save_last=True,
                                          save_top_k=3,
                                          monitor='val/loss',
                                          mode='min',
                                          prefix='')

    if args.restart:
        ckpt = torch.load(glob(
            os.path.join(hparams.log.checkpoint_dir,
                         f'*_epoch={args.resume_from}.ckpt'))[-1],
                          map_location='cpu')
        ckpt_sd = ckpt['state_dict']
        sd = model.state_dict()
        for k, v in sd.items():
            if k in ckpt_sd:
                if ckpt_sd[k].shape == v.shape:
                    sd[k].copy_(ckpt_sd[k])
    if args.ema:
        ckpt = torch.load(glob(
            os.path.join(hparams.log.checkpoint_dir,
                         f'*_epoch={args.resume_from}_EMA'))[-1],
                          map_location='cpu')
        print(ckpt.keys())
        sd = model.state_dict()
        for k, v in sd.items():
            if k in ckpt:
                if ckpt[k].shape == v.shape:
                    sd[k].copy_(ckpt[k])
        args.resume_from = None


    trainer = Trainer(
        checkpoint_callback=checkpoint_callback,
        gpus=hparams.train.gpus,
        accelerator='ddp' if hparams.train.gpus > 1 else None,
        #plugins='ddp_sharded',
        amp_backend='apex',  #
        amp_level='O2',  #
        #num_sanity_val_steps = -1,
        check_val_every_n_epoch=2,
        gradient_clip_val = 1.0,
        max_epochs=200000,
        logger=tblogger,
        progress_bar_refresh_rate=4,
        callbacks=[
            EMACallback(os.path.join(hparams.log.checkpoint_dir,
                         f'{hparams.name}_epoch={{epoch}}_EMA'))
                   ],
        resume_from_checkpoint=None
        if args.resume_from == None or args.restart else sorted(
            glob(
                os.path.join(hparams.log.checkpoint_dir,
                             f'*_epoch={args.resume_from}.ckpt')))[-1])
    trainer.fit(model)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-r', '--resume_from', type =int,\
            required = False, help = "Resume Checkpoint epoch number")
    parser.add_argument('-s', '--restart', action = "store_true",\
            required = False, help = "Significant change occured, use this")
    parser.add_argument('-e', '--ema', action = "store_true",\
            required = False, help = "Start from ema checkpoint")
    args = parser.parse_args()
    #torch.backends.cudnn.benchmark = False
    train(args)