sam_diffsr/tasks/trainer.py

import importlib
import json
import os
import subprocess
import sys
from collections import OrderedDict
from pathlib import Path

parent_path = Path(__file__).absolute().parent.parent
sys.path.append(os.path.abspath(parent_path))
os.chdir(parent_path)
print(f'>-------------> parent path {parent_path}')
print(f'>-------------> current work dir {os.getcwd()}')

cache_path = os.path.join(parent_path, 'cache')
os.environ["HF_DATASETS_CACHE"] = cache_path
os.environ["TRANSFORMERS_CACHE"] = cache_path
os.environ["torch_HOME"] = cache_path

import torch
from PIL import Image
from tqdm import tqdm
import numpy as np
from torch.utils.tensorboard import SummaryWriter
from sam_diffsr.utils_sr.hparams import hparams, set_hparams
from sam_diffsr.utils_sr.utils import plot_img, move_to_cuda, load_checkpoint, save_checkpoint, tensors_to_scalars, Measure, \
    get_all_ckpts



class Trainer:
    def __init__(self):
        self.logger = self.build_tensorboard(save_dir=hparams['work_dir'], name='tb_logs')
        self.measure = Measure()
        self.dataset_cls = None
        self.metric_keys = ['psnr', 'ssim', 'lpips', 'lr_psnr']
        self.metric_2_keys = ['psnr-Y', 'ssim', 'fid']
        self.work_dir = hparams['work_dir']
        self.first_val = True
        
        self.val_steps = hparams['val_steps']
    
    def build_tensorboard(self, save_dir, name, **kwargs):
        log_dir = os.path.join(save_dir, name)
        os.makedirs(log_dir, exist_ok=True)
        return SummaryWriter(log_dir=log_dir, **kwargs)
    
    def build_train_dataloader(self):
        dataset = self.dataset_cls('train')
        return torch.utils.data.DataLoader(
                dataset, batch_size=hparams['batch_size'], shuffle=True,
                pin_memory=False, num_workers=hparams['num_workers'])
    
    def build_val_dataloader(self):
        return torch.utils.data.DataLoader(
                self.dataset_cls('valid'), batch_size=hparams['eval_batch_size'], shuffle=False, pin_memory=False)
    
    def build_test_dataloader(self):
        return torch.utils.data.DataLoader(
                self.dataset_cls('test'), batch_size=hparams['eval_batch_size'], shuffle=False, pin_memory=False)
    
    def build_model(self):
        raise NotImplementedError
    
    def sample_and_test(self, sample):
        raise NotImplementedError
    
    def build_optimizer(self, model):
        raise NotImplementedError
    
    def build_scheduler(self, optimizer):
        raise NotImplementedError
    
    def training_step(self, batch):
        raise NotImplementedError
    
    def train(self):
        model = self.build_model()
        optimizer = self.build_optimizer(model)
        self.global_step = training_step = load_checkpoint(model, optimizer, hparams['work_dir'], steps=self.val_steps)
        self.scheduler = scheduler = self.build_scheduler(optimizer)
        scheduler.step(training_step)
        dataloader = self.build_train_dataloader()
        
        train_pbar = tqdm(dataloader, initial=training_step, total=float('inf'),
                          dynamic_ncols=True, unit='step')
        while self.global_step < hparams['max_updates']:
            for batch in train_pbar:
                if training_step % hparams['val_check_interval'] == 0:
                    with torch.no_grad():
                        model.eval()
                        self.validate(training_step)
                    save_checkpoint(model, optimizer, self.work_dir, training_step, hparams['num_ckpt_keep'])
                model.train()
                batch = move_to_cuda(batch)
                losses, total_loss = self.training_step(batch)
                optimizer.zero_grad()
                
                total_loss.backward()
                optimizer.step()
                training_step += 1
                scheduler.step(training_step)
                self.global_step = training_step
                if training_step % 100 == 0:
                    self.log_metrics({f'tr/{k}': v for k, v in losses.items()}, training_step)
                train_pbar.set_postfix(**tensors_to_scalars(losses))
    
    def validate(self, training_step):
        val_dataloader = self.build_val_dataloader()
        pbar = tqdm(enumerate(val_dataloader), total=len(val_dataloader))
        metrics = {}
        for batch_idx, batch in pbar:
            # 每次运行的第一次validation只跑一小部分数据，来验证代码能否跑通
            if self.first_val and batch_idx > hparams['num_sanity_val_steps'] - 1:
                break
            batch = move_to_cuda(batch)
            img, rrdb_out, ret = self.sample_and_test(batch)
            img_hr = batch['img_hr']
            img_lr = batch['img_lr']
            img_lr_up = batch['img_lr_up']
            if img is not None:
                self.logger.add_image(f'Pred_{batch_idx}', plot_img(img[0]), self.global_step)
                if hparams.get('aux_l1_loss'):
                    self.logger.add_image(f'rrdb_out_{batch_idx}', plot_img(rrdb_out[0]), self.global_step)
                if self.global_step <= hparams['val_check_interval']:
                    self.logger.add_image(f'HR_{batch_idx}', plot_img(img_hr[0]), self.global_step)
                    self.logger.add_image(f'LR_{batch_idx}', plot_img(img_lr[0]), self.global_step)
                    self.logger.add_image(f'BL_{batch_idx}', plot_img(img_lr_up[0]), self.global_step)
            metrics = {}
            metrics.update({k: np.mean(ret[k]) for k in self.metric_keys})
            pbar.set_postfix(**tensors_to_scalars(metrics))
        if hparams['infer']:
            print('Val results:', metrics)
        else:
            if not self.first_val:
                self.log_metrics({f'val/{k}': v for k, v in metrics.items()}, training_step)
                print('Val results:', metrics)
            else:
                print('Sanity val results:', metrics)
        self.first_val = False
    
    def build_test_my_dataloader(self, data_name):
        return torch.utils.data.DataLoader(
                self.dataset_cls(data_name), batch_size=hparams['eval_batch_size'], shuffle=False, pin_memory=False)
    
    def benchmark(self, benchmark_name_list, metric_list):
        from sam_diffsr.tools.caculate_iqa import eval_img_IQA

        model = self.build_model()
        optimizer = self.build_optimizer(model)
        training_step = load_checkpoint(model, optimizer, hparams['work_dir'], hparams['val_steps'])
        self.global_step = training_step
        
        optimizer = None
        
        for data_name in benchmark_name_list:
            test_dataloader = self.build_test_my_dataloader(data_name)
            
            self.results = {k: 0 for k in self.metric_keys}
            self.n_samples = 0
            self.gen_dir = f"{hparams['work_dir']}/results_{self.global_step}_{hparams['gen_dir_name']}/benchmark/{data_name}"
            if hparams['test_save_png']:
                subprocess.check_call(f'rm -rf {self.gen_dir}', shell=True)
                os.makedirs(f'{self.gen_dir}/outputs', exist_ok=True)
                os.makedirs(f'{self.gen_dir}/SR', exist_ok=True)
            
            self.model.sample_tqdm = False
            torch.backends.cudnn.benchmark = False
            if hparams['test_save_png']:
                if hasattr(self.model.denoise_fn, 'make_generation_fast_'):
                    self.model.denoise_fn.make_generation_fast_()
                os.makedirs(f'{self.gen_dir}/HR', exist_ok=True)
            
            result_dict = {}
            
            with torch.no_grad():
                model.eval()
                pbar = tqdm(enumerate(test_dataloader), total=len(test_dataloader))
                for batch_idx, batch in pbar:
                    move_to_cuda(batch)
                    gen_dir = self.gen_dir
                    item_names = batch['item_name']
                    img_hr = batch['img_hr']
                    img_lr = batch['img_lr']
                    img_lr_up = batch['img_lr_up']
                    
                    res = self.sample_and_test(batch)
                    if len(res) == 3:
                        img_sr, rrdb_out, ret = res
                    else:
                        img_sr, ret = res
                        rrdb_out = img_sr
                    
                    img_lr_up = batch.get('img_lr_up', img_lr_up)
                    if img_sr is not None:
                        metrics = list(self.metric_keys)
                        result_dict[batch['item_name'][0]] = {}
                        for k in metrics:
                            self.results[k] += ret[k]
                            result_dict[batch['item_name'][0]][k] = ret[k]
                        self.n_samples += ret['n_samples']
                        
                        print({k: round(self.results[k] / self.n_samples, 3) for k in self.results}, 'total:',
                              self.n_samples)
                        
                        if hparams['test_save_png'] and img_sr is not None:
                            img_sr = self.tensor2img(img_sr)
                            img_hr = self.tensor2img(img_hr)
                            img_lr = self.tensor2img(img_lr)
                            img_lr_up = self.tensor2img(img_lr_up)
                            rrdb_out = self.tensor2img(rrdb_out)
                            for item_name, hr_p, hr_g, lr, lr_up, rrdb_o in zip(
                                    item_names, img_sr, img_hr, img_lr, img_lr_up, rrdb_out):
                                item_name = os.path.splitext(item_name)[0]
                                hr_p = Image.fromarray(hr_p)
                                hr_g = Image.fromarray(hr_g)
                                hr_p.save(f"{gen_dir}/SR/{item_name}.png")
                                hr_g.save(f"{gen_dir}/HR/{item_name}.png")
            
            exp_name = hparams['work_dir'].split('/')[-1]
            sr_img_dir = f"{gen_dir}/SR/"
            gt_img_dir = f"{gen_dir}/HR/"
            excel_path = f"{hparams['work_dir']}/IQA-val-benchmark-{exp_name}.xlsx"
            epoch = training_step
            eval_img_IQA(gt_img_dir, sr_img_dir, excel_path, metric_list, epoch, data_name)
            
            os.makedirs(f'{self.gen_dir}', exist_ok=True)
            eval_json_path = os.path.join(self.gen_dir, 'eval.json')
            avg_result = {k: round(self.results[k] / self.n_samples, 4) for k in self.results}
            with open(eval_json_path, 'w+') as file:
                json.dump(avg_result, file, sort_keys=True, indent=4, separators=(',', ': '), ensure_ascii=False)
                json.dump(result_dict, file, sort_keys=True, indent=4, separators=(',', ': '), ensure_ascii=False)
    
    def benchmark_loop(self, benchmark_name_list, metric_list, gt_path):
        # infer and evaluation all save checkpoint
        from sam_diffsr.tools.caculate_iqa import eval_img_IQA
        
        model = self.build_model()
        
        def get_checkpoint(model, checkpoint):
            stat_dict = checkpoint['state_dict']['model']
            
            new_state_dict = OrderedDict()
            for k, v in stat_dict.items():
                if k[:7] == 'module.':
                    k = k[7:]  # 去掉 `module.`
                new_state_dict[k] = v
            
            model.load_state_dict(new_state_dict)
            model.cuda()
            training_step = checkpoint['global_step']
            del checkpoint
            torch.cuda.empty_cache()
            
            return training_step
        
        ckpt_paths = get_all_ckpts(hparams['work_dir'])
        for ckpt_path in ckpt_paths:
            checkpoint = torch.load(ckpt_path, map_location='cpu')
            training_step = get_checkpoint(model, checkpoint)
            
            self.global_step = training_step
            
            for data_name in benchmark_name_list:
                test_dataloader = self.build_test_my_dataloader(data_name)
                
                self.results = {k: 0 for k in self.metric_keys + self.metric_2_keys}
                self.n_samples = 0
                self.gen_dir = f"{hparams['work_dir']}/results_{training_step}_{hparams['gen_dir_name']}/benchmark/{data_name}"
                
                os.makedirs(f'{self.gen_dir}/outputs', exist_ok=True)
                os.makedirs(f'{self.gen_dir}/SR', exist_ok=True)
                
                self.model.sample_tqdm = False
                torch.backends.cudnn.benchmark = False
                
                with torch.no_grad():
                    model.eval()
                    pbar = tqdm(enumerate(test_dataloader), total=len(test_dataloader))
                    for batch_idx, batch in pbar:
                        move_to_cuda(batch)
                        gen_dir = self.gen_dir
                        item_names = batch['item_name']
                        
                        res = self.sample_and_test(batch)
                        if len(res) == 3:
                            img_sr, rrdb_out, ret = res
                        else:
                            img_sr, ret = res
                            rrdb_out = img_sr
                        
                        img_sr = self.tensor2img(img_sr)
                        
                        for item_name, hr_p in zip(item_names, img_sr):
                            item_name = os.path.splitext(item_name)[0]
                            hr_p = Image.fromarray(hr_p)
                            hr_p.save(f"{gen_dir}/SR/{item_name}.png")
                
                exp_name = hparams['work_dir'].split('/')[-1]
                sr_img_dir = f"{gen_dir}/SR/"
                gt_img_dir = f"{gt_path}/{data_name}/HR"
                excel_path = f"{hparams['work_dir']}/IQA-val-benchmark_loop-{exp_name}.xlsx"
                epoch = training_step
                eval_img_IQA(gt_img_dir, sr_img_dir, excel_path, metric_list, epoch, data_name)
    
    # utils_sr
    def log_metrics(self, metrics, step):
        metrics = self.metrics_to_scalars(metrics)
        logger = self.logger
        for k, v in metrics.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            logger.add_scalar(k, v, step)
    
    def metrics_to_scalars(self, metrics):
        new_metrics = {}
        for k, v in metrics.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            
            if type(v) is dict:
                v = self.metrics_to_scalars(v)
            
            new_metrics[k] = v
        
        return new_metrics
    
    @staticmethod
    def tensor2img(img):
        img = np.round((img.permute(0, 2, 3, 1).cpu().numpy() + 1) * 127.5)
        img = img.clip(min=0, max=255).astype(np.uint8)
        return img


if __name__ == '__main__':
    set_hparams()
    
    pkg = ".".join(hparams["trainer_cls"].split(".")[:-1])
    cls_name = hparams["trainer_cls"].split(".")[-1]
    trainer = getattr(importlib.import_module(pkg), cls_name)()
    if hparams['benchmark_loop']:
        trainer.benchmark_loop(hparams['benchmark_name_list'], hparams['metric_list'], hparams['gt_img_path'])
    elif hparams['benchmark']:
        trainer.benchmark(hparams['benchmark_name_list'], hparams['metric_list'])
    else:
        trainer.train()