Spaces:
Running
Running
| 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 | |
| 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: # separate flow params and normal params for different lr | |
| 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 = [ | |
| { # add normal params first | |
| '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 | |
| # initialize self.metric_results | |
| # It is a dict: { | |
| # 'folder1': tensor (num_frame x len(metrics)), | |
| # 'folder2': tensor (num_frame x len(metrics)) | |
| # } | |
| if with_metrics: | |
| if not hasattr(self, 'metric_results'): # only execute in the first run | |
| 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') | |
| # initialize the best metric results | |
| self._initialize_best_metric_results(dataset_name) | |
| # zero self.metric_results | |
| 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') | |
| # Will evaluate (num_folders + num_pad) times, but only the first num_folders results will be recorded. | |
| # (To avoid wait-dead) | |
| 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'] | |
| # compute outputs | |
| 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() | |
| # tentative for out of GPU memory | |
| 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) | |
| # evaluate | |
| if i < num_folders: | |
| for idx in range(visuals['result'].size(1)): | |
| result = visuals['result'][0, idx, :, :, :] | |
| result_img = tensor2img([result]) # uint8, bgr | |
| metric_data['img'] = result_img | |
| if 'gt' in visuals: | |
| gt = visuals['gt'][0, idx, :, :, :] | |
| gt_img = tensor2img([gt]) # uint8, bgr | |
| 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: # vimeo-90k | |
| 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: # others | |
| img_path = osp.join(self.opt['path']['visualization'], dataset_name, folder, | |
| f"{idx:08d}_{self.opt['name']}.png") | |
| # image name only for REDS dataset | |
| imwrite(result_img, img_path) | |
| # calculate metrics | |
| 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 | |
| # progress bar | |
| 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']: | |
| # collect data among GPUs | |
| 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() | |