| import torch |
| from collections import Counter |
| from os import path as osp |
| from torch import distributed as dist |
| from tqdm import tqdm |
|
|
| from basicsr.metrics import calculate_metric |
| from basicsr.utils import get_root_logger, imwrite, tensor2img |
| from basicsr.utils.dist_util import get_dist_info |
| from basicsr.utils.registry import MODEL_REGISTRY |
| from .video_base_model import VideoBaseModel |
|
|
|
|
| @MODEL_REGISTRY.register() |
| class VideoRecurrentModel(VideoBaseModel): |
|
|
| def __init__(self, opt): |
| super(VideoRecurrentModel, self).__init__(opt) |
| if self.is_train: |
| self.fix_flow_iter = opt['train'].get('fix_flow') |
|
|
| def setup_optimizers(self): |
| train_opt = self.opt['train'] |
| flow_lr_mul = train_opt.get('flow_lr_mul', 1) |
| logger = get_root_logger() |
| logger.info(f'Multiple the learning rate for flow network with {flow_lr_mul}.') |
| if flow_lr_mul == 1: |
| optim_params = self.net_g.parameters() |
| else: |
| normal_params = [] |
| flow_params = [] |
| for name, param in self.net_g.named_parameters(): |
| if 'spynet' in name: |
| flow_params.append(param) |
| else: |
| normal_params.append(param) |
| optim_params = [ |
| { |
| 'params': normal_params, |
| 'lr': train_opt['optim_g']['lr'] |
| }, |
| { |
| 'params': flow_params, |
| 'lr': train_opt['optim_g']['lr'] * flow_lr_mul |
| }, |
| ] |
|
|
| optim_type = train_opt['optim_g'].pop('type') |
| self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g']) |
| self.optimizers.append(self.optimizer_g) |
|
|
| def optimize_parameters(self, current_iter): |
| if self.fix_flow_iter: |
| logger = get_root_logger() |
| if current_iter == 1: |
| logger.info(f'Fix flow network and feature extractor for {self.fix_flow_iter} iters.') |
| for name, param in self.net_g.named_parameters(): |
| if 'spynet' in name or 'edvr' in name: |
| param.requires_grad_(False) |
| elif current_iter == self.fix_flow_iter: |
| logger.warning('Train all the parameters.') |
| self.net_g.requires_grad_(True) |
|
|
| super(VideoRecurrentModel, self).optimize_parameters(current_iter) |
|
|
| def dist_validation(self, dataloader, current_iter, tb_logger, save_img): |
| dataset = dataloader.dataset |
| dataset_name = dataset.opt['name'] |
| with_metrics = self.opt['val']['metrics'] is not None |
| |
| |
| |
| |
| |
| if with_metrics: |
| if not hasattr(self, 'metric_results'): |
| self.metric_results = {} |
| num_frame_each_folder = Counter(dataset.data_info['folder']) |
| for folder, num_frame in num_frame_each_folder.items(): |
| self.metric_results[folder] = torch.zeros( |
| num_frame, len(self.opt['val']['metrics']), dtype=torch.float32, device='cuda') |
| |
| self._initialize_best_metric_results(dataset_name) |
| |
| rank, world_size = get_dist_info() |
| if with_metrics: |
| for _, tensor in self.metric_results.items(): |
| tensor.zero_() |
|
|
| metric_data = dict() |
| num_folders = len(dataset) |
| num_pad = (world_size - (num_folders % world_size)) % world_size |
| if rank == 0: |
| pbar = tqdm(total=len(dataset), unit='folder') |
| |
| |
| for i in range(rank, num_folders + num_pad, world_size): |
| idx = min(i, num_folders - 1) |
| val_data = dataset[idx] |
| folder = val_data['folder'] |
|
|
| |
| val_data['lq'].unsqueeze_(0) |
| val_data['gt'].unsqueeze_(0) |
| self.feed_data(val_data) |
| val_data['lq'].squeeze_(0) |
| val_data['gt'].squeeze_(0) |
|
|
| self.test() |
| visuals = self.get_current_visuals() |
|
|
| |
| del self.lq |
| del self.output |
| if 'gt' in visuals: |
| del self.gt |
| torch.cuda.empty_cache() |
|
|
| if self.center_frame_only: |
| visuals['result'] = visuals['result'].unsqueeze(1) |
| if 'gt' in visuals: |
| visuals['gt'] = visuals['gt'].unsqueeze(1) |
|
|
| |
| if i < num_folders: |
| for idx in range(visuals['result'].size(1)): |
| result = visuals['result'][0, idx, :, :, :] |
| result_img = tensor2img([result]) |
| metric_data['img'] = result_img |
| if 'gt' in visuals: |
| gt = visuals['gt'][0, idx, :, :, :] |
| gt_img = tensor2img([gt]) |
| metric_data['img2'] = gt_img |
|
|
| if save_img: |
| if self.opt['is_train']: |
| raise NotImplementedError('saving image is not supported during training.') |
| else: |
| if self.center_frame_only: |
| clip_ = val_data['lq_path'].split('/')[-3] |
| seq_ = val_data['lq_path'].split('/')[-2] |
| name_ = f'{clip_}_{seq_}' |
| img_path = osp.join(self.opt['path']['visualization'], dataset_name, folder, |
| f"{name_}_{self.opt['name']}.png") |
| else: |
| img_path = osp.join(self.opt['path']['visualization'], dataset_name, folder, |
| f"{idx:08d}_{self.opt['name']}.png") |
| |
| imwrite(result_img, img_path) |
|
|
| |
| if with_metrics: |
| for metric_idx, opt_ in enumerate(self.opt['val']['metrics'].values()): |
| result = calculate_metric(metric_data, opt_) |
| self.metric_results[folder][idx, metric_idx] += result |
|
|
| |
| if rank == 0: |
| for _ in range(world_size): |
| pbar.update(1) |
| pbar.set_description(f'Folder: {folder}') |
|
|
| if rank == 0: |
| pbar.close() |
|
|
| if with_metrics: |
| if self.opt['dist']: |
| |
| for _, tensor in self.metric_results.items(): |
| dist.reduce(tensor, 0) |
| dist.barrier() |
|
|
| if rank == 0: |
| self._log_validation_metric_values(current_iter, dataset_name, tb_logger) |
|
|
| def test(self): |
| n = self.lq.size(1) |
| self.net_g.eval() |
|
|
| flip_seq = self.opt['val'].get('flip_seq', False) |
| self.center_frame_only = self.opt['val'].get('center_frame_only', False) |
|
|
| if flip_seq: |
| self.lq = torch.cat([self.lq, self.lq.flip(1)], dim=1) |
|
|
| with torch.no_grad(): |
| self.output = self.net_g(self.lq) |
|
|
| if flip_seq: |
| output_1 = self.output[:, :n, :, :, :] |
| output_2 = self.output[:, n:, :, :, :].flip(1) |
| self.output = 0.5 * (output_1 + output_2) |
|
|
| if self.center_frame_only: |
| self.output = self.output[:, n // 2, :, :, :] |
|
|
| self.net_g.train() |
|
|