轻量高效！清华智能计算实验室开源基于PyTorch的视频 (图片) 去模糊框架SimDeblur

Python 3 (Conda is recommended)Pytorch 1.5.1 (with GPU)CUDA 10.2+

 git clone https://github.com/ljzycmd/SimDeblur.git

# create a pytorch envconda create -n simdeblur python=3.7conda activate simdeblur   # install the packagescd SimDeblurbash Install.sh

from simdeblur.config import build_config, merge_argsfrom simdeblur.engine.parse_arguments import parse_argumentsfrom simdeblur.engine.trainer import Trainer

args = parse_arguments()
cfg = build_config(args.config_file)cfg = merge_args(cfg, args)cfg.args = args
trainer = Trainer(cfg)trainer.train()

CUDA_VISIBLE_DEVICES=0 bash ./tools/train.sh ./config/dbn/dbn_dvd.yaml 1

CUDA_VISIBLE_DEVICES=0,1,2,3 bash ./tools/train.sh ./config/dbn/dbn_dvd.yaml 4

CONFIG=$1GPUS=$2PORT=${PORT:=10086}# PORT=10086# single gpu trainingif [ GPUS == 1 ]thenecho start single GPU trainingpython train.py $CONFIG --gpus=$GPUSelseecho start distributed training# distributed trainingPYTHONPATH="$(dirname $0)/..":$PYTHONPATH \python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \        train.py $CONFIG --gpus=$GPUSfi

from easydict import EasyDict as edictfrom simdeblur.dataset import build_dataset
dataset = build_dataset(edict({    "name": "DVD",    "mode": "train",    "sampling": "n_c",    "overlapping": True,    "interval": 1,    "root_gt": "./dataset/DVD/quantitative_datasets",    "num_frames": 5,    "augmentation": {        "RandomCrop": {            "size": [256, 256] },        "RandomHorizontalFlip": {            "p": 0.5 },        "RandomVerticalFlip": {            "p": 0.5 },        "RandomRotation90": {            "p": 0.5 },    }}))
print(dataset[0])

from simdeblur.model import build_backbone
model = build_backbone({    "name": "DBN",    "num_frames": 5,    "in_channels": 3,    "inner_channels": 64})
x = torch.randn(1, 5, 3, 256, 256)out = model(x)

from simdeblur.model import build_loss
criterion = build_loss({    "name": "MSELoss",})x = torch.randn(2, 3, 256, 256)y = torch.randn(2, 3, 256, 256)print(criterion(x, y))

import torch
from simdeblur.config import build_config, merge_argsfrom simdeblur.engine.parse_arguments import parse_argumentsfrom simdeblur.engine.trainer import Trainer

def main():    args = parse_arguments()
    cfg = build_config(args.config_file)    cfg = merge_args(cfg, args)    cfg.args = args
    trainer = Trainer(cfg)    trainer.train()    
if __name__ == "__main__":    main()

from simdeblur.dataset import build_datasetfrom simdeblur.scheduler import build_optimizer, build_lr_schedulerfrom simdeblur.model import build_backbone, build_meta_arch, build_lossfrom simdeblur.utils.logger import LogBuffer, SimpleMetricPrinter, TensorboardWriterfrom simdeblur.utils.metrics import calculate_psnr, calculate_ssimfrom simdeblur.utils import dist_utils
from simdeblur.engine import hooks

logging.basicConfig(format='%(asctime)s - %(levelname)s - SimDeblur: %(message)s',level=logging.INFO)logging.info("******* A simple deblurring framework ********")
class Trainer:    def __init__(self, cfg):        """        Args            cfg(edict): the config file, which contains arguments form comand line        """        self.cfg = copy.deepcopy(cfg)        # initialize the distributed training        if cfg.args.gpus > 1:            dist_utils.init_distributed(cfg)
        # create the working dirs        self.current_work_dir = os.path.join(cfg.work_dir, cfg.name)        if not os.path.exists(self.current_work_dir):            os.makedirs(self.current_work_dir, exist_ok=True)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")        # self.device = torch.device("cpu")
        # default logger        logger = logging.getLogger("simdeblur")        logger.setLevel(logging.INFO)        logger.addHandler(            logging.FileHandler(                os.path.join(                    self.current_work_dir, self.cfg.name.split("_")[0] + ".json"))        )                # construct the modules        self.model = self.build_model(cfg).to(self.device)        self.criterion = build_loss(cfg.loss).to(self.device)        self.train_dataloader, self.train_sampler = self.build_dataloder(cfg, mode="train")        self.val_datalocaer, _ = self.build_dataloder(cfg, mode="val")        self.optimizer = self.build_optimizer(cfg, self.model)        self.lr_scheduler = self.build_lr_scheduler(cfg, self.optimizer)                # trainer hooks        self._hooks = self.build_hooks()
        # some induces when training        self.epochs = 0        self.iters = 0        self.batch_idx = 0 
        self.start_epoch = 0        self.start_iter = 0        self.total_train_epochs = self.cfg.schedule.epochs        self.total_train_iters = self.total_train_epochs * len(self.train_dataloader)
        # resume or load the ckpt as init-weights        if self.cfg.resume_from != "None":            self.resume_or_load_ckpt(ckpt_path=self.cfg.resume_from)
        # log bufffer(dict to save)         self.log_buffer = LogBuffer()

    def before_epoch(self):        for h in self._hooks:            h.before_epoch(self)        # shuffle the data when dist training ...        if self.train_sampler:            self.train_sampler.set_epoch(self.epochs)

    def before_epoch(self):        for h in self._hooks:            h.before_epoch(self)        # shuffle the data when dist training ...        if self.train_sampler:            self.train_sampler.set_epoch(self.epochs)

    def preprocess(self, batch_data):        """        prepare for input        """        return batch_data["input_frames"].to(self.device)

 def postprocess(self):        """        post process for model outputs        """        # When the outputs is a img tensor        if isinstance(self.outputs, torch.Tensor) and self.outputs.dim() == 5:            self.outputs = self.outputs.flatten(0, 1)

    def calculate_loss(self, batch_data, model_outputs):        """        calculate the loss        """        gt_frames = batch_data["gt_frames"].to(self.device).flatten(0, 1)        if model_outputs.dim() == 5:                model_outputs = model_outputs.flatten(0, 1) # (b*n, c, h, w)        return self.criterion(gt_frames, model_outputs)

    def update_params(self):        """        update params        pipline: zero_grad, backward and update grad        """        self.optimizer.zero_grad()        self.loss.backward()        self.optimizer.step()

    def after_iter(self):        for h in self._hooks:            h.after_iter(self)    def after_epoch(self):        for h in self._hooks:            h.after_epoch(self)

    def train(self, **kwargs):        self.model.train()        self.before_train()        logger = logging.getLogger("simdeblur")        logger.info("Starting training...")        for self.epochs in range(self.start_epoch, self.cfg.schedule.epochs):            # shuffle the dataloader when dist training: dist_data_loader.set_epoch(epoch)            self.before_epoch()            for self.batch_idx, self.batch_data in enumerate(self.train_dataloader):                self.before_iter()
                input_frames = self.preprocess(self.batch_data)
                self.outputs = self.model(input_frames)                self.postprocess()
                self.loss = self.calculate_loss(self.batch_data, self.outputs)
                self.update_params()
                self.iters += 1                self.after_iter()                        if self.epochs % self.cfg.schedule.val_epochs == 0:                self.val()
            self.after_epoch()

@Article{wang2021simdeblur,  author  = {Mingdeng Cao, Jiahao Wang},  title   = {清华智能计算实验室团队开源基于PyTorch的视频 (图片) 去模糊框架SimDeblur},  journal = {https://zhuanlan.zhihu.com/},  howpublished = {\url{https://github.com/ljzycmd/SimDeblur}},  year    = {2021},  url= {https://zhuanlan.zhihu.com/p/368312516/},}