devbox
运维做开发
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

超分辨率(3)--基于RCAN实现图像超分辨率重建_rcan网络

作者：运维做开发 | 2024-07-03 10:30:21

赞

踩

rcan网络

目录

一.项目介绍

二.项目流程详解

2.1.数据处理模块

2.2.损失函数设置 

2.3.网络模型构建

三.测试网络

一.项目介绍

RCAN：Residual Channel Attention Network（残差通道注意网络 )

卷积神经网络(CNN)的深度对于图像超分辨率(SR)是极其关键的因素。然而，我们观察到，更深层次的图像SR网络更难训练。低分辨率的输入和特征包含丰富的低频信息，这些信息在通道间被平等对待，从而阻碍了CNNs的表征能力。为了解决这些问题，我们提出了一种非常深的残差通道注意网络(RCAN)。具体地，我们提出了一种residual in residual（RIR）结构来形成非常深的网络，它由几个具有长跳连接的残差组组成。每个残差组包含一些具有短跳连接的残差块。与此同时，RIR允许大量的低频信息通过多个跳跃连接被绕过，使得主网络专注于学习高频信息。在此基础上，我们提出了一种通道注意机制，通过考虑通道间的相互依赖关系，自适应地重新调整通道特征。大量的实验表明，与比之前最先进的方法相比，我们的RCAN实现了更好的精度和视觉效果。

背景：

  • 卷积神经网络（CNN）的深度对于图像超分辨率（SR）是极其关键的因素。然而，作者观察到，更深层次的图像SR网络更难训练。
  • 低分辨率图像（LR）的输入和特征包含大量的低频信息，这些信息在通道间被平等对待，从而阻碍了CNNs的表征能力。

解决方案：

  • 对于第一个更深的网络更难训练的问题，作者研究发现，通过在网络中引入残差块，这种残差块使得网络达到了1000层，但是仅仅通过叠加残差块来构建更深的网络很难获得更好的提升效果。因此，作者提出了残差嵌套（residual in residual，RIR）结构构造非常深的可训练网络，RIR中的长跳连接和短跳连接有助于绕过大量的低频信息，使主网络学习到更有效的信息。
  • 对于第二个LR输入低频和高频信息在通道被平等对待的问题，作者发现注意力可以使可用处理资源的分配偏向于输入中信息量最大的部分，因此引入通道注意（Channel Attention ，CA）机制。

网络架构：

​

​

​

RCAN主要由四个部分组成：浅层特征提取、残差嵌套（RIR）深度特征提取、上采样模块和重建部分。 

  •  RIR组成：G个RG（带长跳连接）
  • 每个RG：B个RCAB组成（带短跳连接）
  • 每个RCAB组成：Conv + ReLU + Conv + CA
  • CA组成：Global pooling + Conv + ReLU + Conv  

名词解释：

  • Residual Channel Attention Network，RCAN：残差通道注意网络
  • residual in residua，RIR：残差嵌套
  • residual groups，RG：残差组
  • Residual Channel Attention Block，RCAB：残差通道注意块
  • Channel Attention，CA：通道注意
  • long skip connection，LSC：长跳连接
  • short skip connection，SSC：短跳连接

论文地址：

[1807.02758] Image Super-Resolution Using Very Deep Residual Channel Attention Networks (arxiv.org)icon-default.png?t=N7T8https://arxiv.org/abs/1807.02758

参考文章： 

RCAN论文笔记：Image Super-Resolution Using Very Deep Residual Channel Attention Networks-CSDN博客icon-default.png?t=N7T8https://blog.csdn.net/weixin_46773169/article/details/105600346

源码地址：

yulunzhang/RCAN: PyTorch code for our ECCV 2018 paper "Image Super-Resolution Using Very Deep Residual Channel Attention Networks" (github.com)icon-default.png?t=N7T8https://github.com/yulunzhang/RCAN

二.项目流程详解

2.1.数据处理模块

_init_.py

  1. from importlib import import_module
  2.  
  3. from dataloader import MSDataLoader
  4. from torch.utils.data.dataloader import default_collate
  5.  
  6. class Data:
  7.     def __init__(self, args):
  8.         kwargs = {}
  9.         # 如果不在cpu上训练
  10.         if not args.cpu:
  11.             kwargs['collate_fn'] = default_collate
  12.             kwargs['pin_memory'] = True
  13.         # 在cpu上训练
  14.         else:
  15.             kwargs['collate_fn'] = default_collate
  16.             kwargs['pin_memory'] = False
  17.  
  18.         self.loader_train = None
  19.         if not args.test_only:
  20.             # .lower()将大写字母转换为小写字母
  21.             module_train = import_module('data.' + args.data_train.lower())
  22.             # getattr() 函数用于返回一个对象属性值。
  23.             trainset = getattr(module_train, args.data_train)(args)
  24.             self.loader_train = MSDataLoader(
  25.                 args,
  26.                 trainset,
  27.                 batch_size=args.batch_size,
  28.                 shuffle=True,
  29.                 **kwargs
  30.             )
  31.  
  32.         # 针对特殊的数据
  33.         if args.data_test in ['Set5', 'Set14', 'B100', 'Urban100']:
  34.             if not args.benchmark_noise:
  35.                 module_test = import_module('data.benchmark')
  36.                 testset = getattr(module_test, 'Benchmark')(args, train=False)
  37.             else:
  38.                 module_test = import_module('data.benchmark_noise')
  39.                 testset = getattr(module_test, 'BenchmarkNoise')(
  40.                     args,
  41.                     train=False
  42.                 )
  43.  
  44.         else:
  45.             module_test = import_module('data.' +  args.data_test.lower())
  46.             testset = getattr(module_test, args.data_test)(args, train=False)
  47.  
  48.         # 对于自定义的MSDataLoader，主要需要传入的参数为args和dataset
  49.         self.loader_test = MSDataLoader(
  50.             args,
  51.             testset,
  52.             batch_size=1,
  53.             shuffle=False,
  54.             **kwargs
  55.         )
  56.  
  57. '''
  58. class MSDataLoader(DataLoader):
  59.     def __init__(
  60.         self, args, dataset, batch_size=1, shuffle=False,
  61.         sampler=None, batch_sampler=None,
  62.         collate_fn=default_collate, pin_memory=False, drop_last=False,
  63.         timeout=0, worker_init_fn=None):
  65.         super(MSDataLoader, self).__init__(
  66.             dataset, batch_size=batch_size, shuffle=shuffle,
  67.             sampler=sampler, batch_sampler=batch_sampler,
  68.             num_workers=args.n_threads, collate_fn=collate_fn,
  69.             pin_memory=pin_memory, drop_last=drop_last,
  70.             timeout=timeout, worker_init_fn=worker_init_fn)
  72.         self.scale = args.scale
  74.     def __iter__(self):
  75.         return _MSDataLoaderIter(self)
  76. '''

benchmark.py

  1. import os
  2.  
  3. from data import common
  4. from data import srdata
  5.  
  6. import numpy as np
  7. import scipy.misc as misc
  8.  
  9. import torch
  10. import torch.utils.data as data
  11.  
  12. class Benchmark(srdata.SRData):
  13.     def __init__(self, args, train=True):
  14.         super(Benchmark, self).__init__(args, train, benchmark=True)
  15.  
  16.     # 扫描磁盘得到数据
  17.     def _scan(self):
  18.         list_hr = []
  19.         list_lr = [[] for _ in self.scale]
  20.         for entry in os.scandir(self.dir_hr):
  21.             # os.path.splitext分离文件名字和文件类型
  22.             # eg: os.path.splitext(abc.txt) 得到的为('abc','txt')
  23.             # filename取出的是文件名
  24.             filename = os.path.splitext(entry.name)[0]
  25.             # filename + self.ext 为文件的完整名字
  26.             # os.path.join用于拼接文件路径，可以传入多个路径
  27.             # 此处append的文件路径即为 self.dir_hr + (filename+self.ext)
  28.             list_hr.append(os.path.join(self.dir_hr, filename + self.ext))
  29.             for si, s in enumerate(self.scale):
  30.                 list_lr[si].append(os.path.join(
  31.                     self.dir_lr,
  32.                     'X{}/{}x{}{}'.format(s, filename, s, self.ext)
  33.                 ))
  34.  
  35.         # 对取出的数据进行升序排列
  36.         list_hr.sort()
  37.         for l in list_lr:
  38.             l.sort()
  39.  
  40.         return list_hr, list_lr
  41.  
  42.     # 设置数据的地址以及数据的类型
  43.     def _set_filesystem(self, dir_data):
  44.         self.apath = os.path.join(dir_data, 'benchmark', self.args.data_test)
  45.         self.dir_hr = os.path.join(self.apath, 'HR')
  46.         self.dir_lr = os.path.join(self.apath, 'LR_bicubic')
  47.         self.ext = '.png'

 common.py

  1. import random
  2.  
  3. import numpy as np
  4. import skimage.io as sio
  5. import skimage.color as sc
  6. import skimage.transform as st
  7.  
  8. import torch
  9. from torchvision import transforms
  10.  
  11. def get_patch(img_in, img_tar, patch_size, scale, multi_scale=False):
  12.     # shape得到图片的高度、宽度、颜色通道
  13.     # 所以shape[:2}就是获取图片的前两个维度，获得图片的高度和宽度
  14.  
  15.     ih, iw = img_in.shape[:2]
  16.  
  17.     p = scale if multi_scale else 1
  18.     tp = p * patch_size
  19.     ip = tp // scale
  20.  
  21.     ix = random.randrange(0, iw - ip + 1)
  22.     iy = random.randrange(0, ih - ip + 1)
  23.     tx, ty = scale * ix, scale * iy
  24.  
  25.     img_in = img_in[iy:iy + ip, ix:ix + ip, :]
  26.     img_tar = img_tar[ty:ty + tp, tx:tx + tp, :]
  27.  
  28.     return img_in, img_tar
  29.  
  30. # 设置channel值
  31. def set_channel(l, n_channel):
  32.     def _set_channel(img):
  33.         if img.ndim == 2:
  34.             # expand_dims(a, axis)中，a为numpy数组，axis为需添加维度的轴
  35.             # 使数据增加一个维度
  36.             img = np.expand_dims(img, axis=2)
  37.  
  38.         c = img.shape[2]
  39.         if n_channel == 1 and c == 3:
  40.             img = np.expand_dims(sc.rgb2ycbcr(img)[:, :, 0], 2)
  41.         elif n_channel == 3 and c == 1:
  42.             # numpy.concatenate((a1,a2,...), axis=0)函数。
  43.             # 能 够一次完成多个数组的拼接。其中a1,a2,...是数组类型的参数
  44.             img = np.concatenate([img] * n_channel, 2)
  45.  
  46.         return img
  47.  
  48.     return [_set_channel(_l) for _l in l]
  49.  
  50. # 将np.array类型转为tensor类型
  51. def np2Tensor(l, rgb_range):
  52.     def _np2Tensor(img):
  53.         # ascontiguousarray函数将一个内存不连续存储的数组转换为内存连续存储的数组，使得运行速度更快
  54.         # img.transpose((2,0,1))将图片的维度由(0,1,2)转换为(2,0,1)
  55.         np_transpose = np.ascontiguousarray(img.transpose((2, 0, 1)))
  56.         tensor = torch.from_numpy(np_transpose).float()
  57.         tensor.mul_(rgb_range / 255)
  58.  
  59.         return tensor
  60.  
  61.     return [_np2Tensor(_l) for _l in l]
  62.  
  63. def add_noise(x, noise='.'):
  64.     if noise is not '.':
  65.         noise_type = noise[0]
  66.         noise_value = int(noise[1:])
  67.         if noise_type == 'G':
  68.             noises = np.random.normal(scale=noise_value, size=x.shape)
  69.             noises = noises.round()
  70.         elif noise_type == 'S':
  71.             noises = np.random.poisson(x * noise_value) / noise_value
  72.             noises = noises - noises.mean(axis=0).mean(axis=0)
  73.  
  74.         x_noise = x.astype(np.int16) + noises.astype(np.int16)
  75.         x_noise = x_noise.clip(0, 255).astype(np.uint8)
  76.         return x_noise
  77.     else:
  78.         return x
  79.  
  80. def augment(l, hflip=True, rot=True):
  81.     hflip = hflip and random.random() < 0.5
  82.     vflip = rot and random.random() < 0.5
  83.     rot90 = rot and random.random() < 0.5
  84.  
  85.     def _augment(img):
  86.         if hflip: img = img[:, ::-1, :]
  87.         if vflip: img = img[::-1, :, :]
  88.         if rot90: img = img.transpose(1, 0, 2)
  89.         
  90.         return img
  91.  
  92.     return [_augment(_l) for _l in l]

demo.py 

  1. import os
  2.  
  3. from data import common
  4.  
  5. import numpy as np
  6. import scipy.misc as misc
  7.  
  8. import torch
  9. import torch.utils.data as data
  10.  
  11. class Demo(data.Dataset):
  12.     def __init__(self, args, train=False):
  13.         self.args = args
  14.         self.name = 'Demo'
  15.         self.scale = args.scale
  16.         self.idx_scale = 0
  17.         self.train = False
  18.         self.benchmark = False
  19.  
  20.         self.filelist = []
  21.         for f in os.listdir(args.dir_demo):
  22.             if f.find('.png') >= 0 or f.find('.jp') >= 0:
  23.                 self.filelist.append(os.path.join(args.dir_demo, f))
  24.         self.filelist.sort()
  25.  
  26.     def __getitem__(self, idx):
  27.         filename = os.path.split(self.filelist[idx])[-1]
  28.         filename, _ = os.path.splitext(filename)
  29.         lr = misc.imread(self.filelist[idx])
  30.         lr = common.set_channel([lr], self.args.n_colors)[0]
  31.  
  32.         return common.np2Tensor([lr], self.args.rgb_range)[0], -1, filename
  33.  
  34.     def __len__(self):
  35.         return len(self.filelist)
  36.  
  37.     def set_scale(self, idx_scale):
  38.         self.idx_scale = idx_scale
  39.

srdata.py

  1. import os
  2.  
  3. from data import common
  4.  
  5. import numpy as np
  6. import scipy.misc as misc
  7.  
  8. import torch
  9. import torch.utils.data as data
  10.  
  11. class SRData(data.Dataset):
  12.     def __init__(self, args, train=True, benchmark=False):
  13.         self.args = args
  14.         self.train = train
  15.         self.split = 'train' if train else 'test'
  16.         self.benchmark = benchmark
  17.         self.scale = args.scale
  18.         self.idx_scale = 0
  19.  
  20.         self._set_filesystem(args.dir_data)
  21.  
  22.         def _load_bin():
  23.             self.images_hr = np.load(self._name_hrbin())
  24.             self.images_lr = [
  25.                 np.load(self._name_lrbin(s)) for s in self.scale
  26.             ]
  27.  
  28.         if args.ext == 'img' or benchmark:
  29.             self.images_hr, self.images_lr = self._scan()
  30.         elif args.ext.find('sep') >= 0:
  31.             self.images_hr, self.images_lr = self._scan()
  32.             if args.ext.find('reset') >= 0:
  33.                 print('Preparing seperated binary files')
  34.                 for v in self.images_hr:
  35.                     hr = misc.imread(v)
  36.                     name_sep = v.replace(self.ext, '.npy')
  37.                     np.save(name_sep, hr)
  38.                 for si, s in enumerate(self.scale):
  39.                     for v in self.images_lr[si]:
  40.                         lr = misc.imread(v)
  41.                         name_sep = v.replace(self.ext, '.npy')
  42.                         np.save(name_sep, lr)
  43.  
  44.             self.images_hr = [
  45.                 v.replace(self.ext, '.npy') for v in self.images_hr
  46.             ]
  47.             self.images_lr = [
  48.                 [v.replace(self.ext, '.npy') for v in self.images_lr[i]]
  49.                 for i in range(len(self.scale))
  50.             ]
  51.  
  52.         elif args.ext.find('bin') >= 0:
  53.             try:
  54.                 if args.ext.find('reset') >= 0:
  55.                     raise IOError
  56.                 print('Loading a binary file')
  57.                 _load_bin()
  58.             except:
  59.                 print('Preparing a binary file')
  60.                 bin_path = os.path.join(self.apath, 'bin')
  61.                 if not os.path.isdir(bin_path):
  62.                     os.mkdir(bin_path)
  63.  
  64.                 list_hr, list_lr = self._scan()
  65.                 hr = [misc.imread(f) for f in list_hr]
  66.                 np.save(self._name_hrbin(), hr)
  67.                 del hr
  68.                 for si, s in enumerate(self.scale):
  69.                     lr_scale = [misc.imread(f) for f in list_lr[si]]
  70.                     np.save(self._name_lrbin(s), lr_scale)
  71.                     del lr_scale
  72.                 _load_bin()
  73.         else:
  74.             print('Please define data type')
  75.  
  76.     def _scan(self):
  77.         raise NotImplementedError
  78.  
  79.     def _set_filesystem(self, dir_data):
  80.         raise NotImplementedError
  81.  
  82.     def _name_hrbin(self):
  83.         raise NotImplementedError
  84.  
  85.     def _name_lrbin(self, scale):
  86.         raise NotImplementedError
  87.  
  88.     def __getitem__(self, idx):
  89.         lr, hr, filename = self._load_file(idx)
  90.         lr, hr = self._get_patch(lr, hr)
  91.         lr, hr = common.set_channel([lr, hr], self.args.n_colors)
  92.         lr_tensor, hr_tensor = common.np2Tensor([lr, hr], self.args.rgb_range)
  93.         return lr_tensor, hr_tensor, filename
  94.  
  95.     def __len__(self):
  96.         return len(self.images_hr)
  97.  
  98.     def _get_index(self, idx):
  99.         return idx
  100.  
  101.     def _load_file(self, idx):
  102.         idx = self._get_index(idx)
  103.         lr = self.images_lr[self.idx_scale][idx]
  104.         hr = self.images_hr[idx]
  105.         if self.args.ext == 'img' or self.benchmark:
  106.             filename = hr
  107.             lr = misc.imread(lr)
  108.             hr = misc.imread(hr)
  109.         elif self.args.ext.find('sep') >= 0:
  110.             filename = hr
  111.             lr = np.load(lr)
  112.             hr = np.load(hr)
  113.         else:
  114.             filename = str(idx + 1)
  115.  
  116.         filename = os.path.splitext(os.path.split(filename)[-1])[0]
  117.  
  118.         return lr, hr, filename
  119.  
  120.     def _get_patch(self, lr, hr):
  121.         patch_size = self.args.patch_size
  122.         scale = self.scale[self.idx_scale]
  123.         multi_scale = len(self.scale) > 1
  124.         if self.train:
  125.             lr, hr = common.get_patch(
  126.                 lr, hr, patch_size, scale, multi_scale=multi_scale
  127.             )
  128.             lr, hr = common.augment([lr, hr])
  129.             lr = common.add_noise(lr, self.args.noise)
  130.         else:
  131.             ih, iw = lr.shape[0:2]
  132.             hr = hr[0:ih * scale, 0:iw * scale]
  133.  
  134.         return lr, hr
  135.  
  136.     def set_scale(self, idx_scale):
  137.         self.idx_scale = idx_scale
  138.

div2k.py 

  1. import os
  2.  
  3. from data import common
  4. from data import srdata
  5.  
  6. import numpy as np
  7. import scipy.misc as misc
  8.  
  9. import torch
  10. import torch.utils.data as data
  11.  
  12. class DIV2K(srdata.SRData):
  13.     def __init__(self, args, train=True):
  14.         super(DIV2K, self).__init__(args, train)
  15.         self.repeat = args.test_every // (args.n_train // args.batch_size)
  16.  
  17.     def _scan(self):
  18.         list_hr = []
  19.         list_lr = [[] for _ in self.scale]
  20.         if self.train:
  21.             idx_begin = 0
  22.             idx_end = self.args.n_train
  23.         else:
  24.             idx_begin = self.args.n_train
  25.             idx_end = self.args.offset_val + self.args.n_val
  26.  
  27.         for i in range(idx_begin + 1, idx_end + 1):
  28.             filename = '{:0>4}'.format(i)
  29.             list_hr.append(os.path.join(self.dir_hr, filename + self.ext))
  30.             for si, s in enumerate(self.scale):
  31.                 list_lr[si].append(os.path.join(
  32.                     self.dir_lr,
  33.                     'X{}/{}x{}{}'.format(s, filename, s, self.ext)
  34.                 ))
  35.  
  36.         return list_hr, list_lr
  37.  
  38.     def _set_filesystem(self, dir_data):
  39.         self.apath = dir_data + '/DIV2K'
  40.         self.dir_hr = os.path.join(self.apath, 'DIV2K_train_HR')
  41.         self.dir_lr = os.path.join(self.apath, 'DIV2K_train_LR_bicubic')
  42.         self.ext = '.png'
  43.  
  44.     def _name_hrbin(self):
  45.         return os.path.join(
  46.             self.apath,
  47.             'bin',
  48.             '{}_bin_HR.npy'.format(self.split)
  49.         )
  50.  
  51.     def _name_lrbin(self, scale):
  52.         return os.path.join(
  53.             self.apath,
  54.             'bin',
  55.             '{}_bin_LR_X{}.npy'.format(self.split, scale)
  56.         )
  57.  
  58.     def __len__(self):
  59.         if self.train:
  60.             return len(self.images_hr) * self.repeat
  61.         else:
  62.             return len(self.images_hr)
  63.  
  64.     def _get_index(self, idx):
  65.         if self.train:
  66.             return idx % len(self.images_hr)
  67.         else:
  68.             return idx
  69.

2.2.损失函数设置 

_init_.py

  1. import os
  2. from importlib import import_module
  3.  
  4. import matplotlib
  5. matplotlib.use('Agg')
  6. import matplotlib.pyplot as plt
  7.  
  8. import numpy as np
  9.  
  10. import torch
  11. import torch.nn as nn
  12. import torch.nn.functional as F
  13.  
  14. class Loss(nn.modules.loss._Loss):
  15.     def __init__(self, args, ckp):
  16.         super(Loss, self).__init__()
  17.         print('Preparing loss function:')
  18.  
  19.         self.n_GPUs = args.n_GPUs
  20.         self.loss = []
  21.         # 首先说说 nn.ModuleList 这个类，你可以把任意 nn.Module 的子类
  22.         # (比如 nn.Conv2d, nn.Linear 之类的) 加到这个 list 里面，
  23.         # 方法和 Python 自带的 list 一样，无非是 extend，append 等操作。
  24.         # 但不同于一般的 list，加入到 nn.ModuleList 里面的 module 是会自动注册到整个网络上的，
  25.         # 同时 module 的 parameters 也会自动添加到整个网络中。
  26.         self.loss_module = nn.ModuleList()
  27.         # split(' ')根据括号里的字符分割字符串
  28.         for loss in args.loss.split('+'):
  29.             weight, loss_type = loss.split('*')
  30.             if loss_type == 'MSE':
  31.                 loss_function = nn.MSELoss()
  32.             elif loss_type == 'L1':
  33.                 loss_function = nn.L1Loss()
  34.             elif loss_type.find('VGG') >= 0:
  35.                 module = import_module('loss.vgg')
  36.                 loss_function = getattr(module, 'VGG')(
  37.                     loss_type[3:],
  38.                     rgb_range=args.rgb_range
  39.                 )
  40.             elif loss_type.find('GAN') >= 0:
  41.                 module = import_module('loss.adversarial')
  42.                 loss_function = getattr(module, 'Adversarial')(
  43.                     args,
  44.                     loss_type
  45.                 )
  46.            
  47.             self.loss.append({
  48.                 'type': loss_type,
  49.                 'weight': float(weight),
  50.                 'function': loss_function}
  51.             )
  52.             if loss_type.find('GAN') >= 0:
  53.                 self.loss.append({'type': 'DIS', 'weight': 1, 'function': None})
  54.  
  55.         if len(self.loss) > 1:
  56.             self.loss.append({'type': 'Total', 'weight': 0, 'function': None})
  57.  
  58.         for l in self.loss:
  59.             if l['function'] is not None:
  60.                 print('{:.3f} * {}'.format(l['weight'], l['type']))
  61.                 self.loss_module.append(l['function'])
  62.  
  63.         self.log = torch.Tensor()
  64.  
  65.         device = torch.device('cpu' if args.cpu else 'cuda')
  66.         self.loss_module.to(device)
  67.         if args.precision == 'half': self.loss_module.half()
  68.         if not args.cpu and args.n_GPUs > 1:
  69.             self.loss_module = nn.DataParallel(
  70.                 self.loss_module, range(args.n_GPUs)
  71.             )
  72.  
  73.         if args.load != '.': self.load(ckp.dir, cpu=args.cpu)
  74.  
  75.     def forward(self, sr, hr):
  76.         losses = []
  77.         for i, l in enumerate(self.loss):
  78.             if l['function'] is not None:
  79.                 loss = l['function'](sr, hr)
  80.                 effective_loss = l['weight'] * loss
  81.                 losses.append(effective_loss)
  82.                 self.log[-1, i] += effective_loss.item()
  83.             elif l['type'] == 'DIS':
  84.                 self.log[-1, i] += self.loss[i - 1]['function'].loss
  85.  
  86.         loss_sum = sum(losses)
  87.         if len(self.loss) > 1:
  88.             self.log[-1, -1] += loss_sum.item()
  89.  
  90.         return loss_sum
  91.  
  92.     def step(self):
  93.         for l in self.get_loss_module():
  94.             if hasattr(l, 'scheduler'):
  95.                 l.scheduler.step()
  96.  
  97.     def start_log(self):
  98.         self.log = torch.cat((self.log, torch.zeros(1, len(self.loss))))
  99.  
  100.     def end_log(self, n_batches):
  101.         self.log[-1].div_(n_batches)
  102.  
  103.     def display_loss(self, batch):
  104.         n_samples = batch + 1
  105.         log = []
  106.         for l, c in zip(self.loss, self.log[-1]):
  107.             log.append('[{}: {:.4f}]'.format(l['type'], c / n_samples))
  108.  
  109.         return ''.join(log)
  110.  
  111.     def plot_loss(self, apath, epoch):
  112.         axis = np.linspace(1, epoch, epoch)
  113.         for i, l in enumerate(self.loss):
  114.             label = '{} Loss'.format(l['type'])
  115.             fig = plt.figure()
  116.             plt.title(label)
  117.             plt.plot(axis, self.log[:, i].numpy(), label=label)
  118.             plt.legend()
  119.             plt.xlabel('Epochs')
  120.             plt.ylabel('Loss')
  121.             plt.grid(True)
  122.             plt.savefig('{}/loss_{}.pdf'.format(apath, l['type']))
  123.             plt.close(fig)
  124.  
  125.     def get_loss_module(self):
  126.         if self.n_GPUs == 1:
  127.             return self.loss_module
  128.         else:
  129.             return self.loss_module.module
  130.  
  131.     def save(self, apath):
  132.         torch.save(self.state_dict(), os.path.join(apath, 'loss.pt'))
  133.         torch.save(self.log, os.path.join(apath, 'loss_log.pt'))
  134.  
  135.     def load(self, apath, cpu=False):
  136.         if cpu:
  137.             kwargs = {'map_location': lambda storage, loc: storage}
  138.         else:
  139.             kwargs = {}
  140.  
  141.         self.load_state_dict(torch.load(
  142.             os.path.join(apath, 'loss.pt'),
  143.             **kwargs
  144.         ))
  145.         self.log = torch.load(os.path.join(apath, 'loss_log.pt'))
  146.         for l in self.get_loss_module():
  147.             if hasattr(l, 'scheduler'):
  148.                 for _ in range(len(self.log)): l.scheduler.step()
  149.

 adversarial.py

  1. import utility
  2. from model import common
  3. from loss import discriminator
  4.  
  5. import torch
  6. import torch.nn as nn
  7. import torch.nn.functional as F
  8. import torch.optim as optim
  9. from torch.autograd import Variable
  10.  
  11. class Adversarial(nn.Module):
  12.     def __init__(self, args, gan_type):
  13.         super(Adversarial, self).__init__()
  14.         self.gan_type = gan_type
  15.         self.gan_k = args.gan_k
  16.         self.discriminator = discriminator.Discriminator(args, gan_type)
  17.         if gan_type != 'WGAN_GP':
  18.             self.optimizer = utility.make_optimizer(args, self.discriminator)
  19.         else:
  20.             self.optimizer = optim.Adam(
  21.                 self.discriminator.parameters(),
  22.                 betas=(0, 0.9), eps=1e-8, lr=1e-5
  23.             )
  24.         self.scheduler = utility.make_scheduler(args, self.optimizer)
  25.  
  26.     def forward(self, fake, real):
  27.         fake_detach = fake.detach()
  28.  
  29.         self.loss = 0
  30.         for _ in range(self.gan_k):
  31.             self.optimizer.zero_grad()
  32.             d_fake = self.discriminator(fake_detach)
  33.             d_real = self.discriminator(real)
  34.             if self.gan_type == 'GAN':
  35.                 label_fake = torch.zeros_like(d_fake)
  36.                 label_real = torch.ones_like(d_real)
  37.                 loss_d \
  38.                     = F.binary_cross_entropy_with_logits(d_fake, label_fake) \
  39.                     + F.binary_cross_entropy_with_logits(d_real, label_real)
  40.             elif self.gan_type.find('WGAN') >= 0:
  41.                 loss_d = (d_fake - d_real).mean()
  42.                 if self.gan_type.find('GP') >= 0:
  43.                     epsilon = torch.rand_like(fake).view(-1, 1, 1, 1)
  44.                     hat = fake_detach.mul(1 - epsilon) + real.mul(epsilon)
  45.                     hat.requires_grad = True
  46.                     d_hat = self.discriminator(hat)
  47.                     gradients = torch.autograd.grad(
  48.                         outputs=d_hat.sum(), inputs=hat,
  49.                         retain_graph=True, create_graph=True, only_inputs=True
  50.                     )[0]
  51.                     gradients = gradients.view(gradients.size(0), -1)
  52.                     gradient_norm = gradients.norm(2, dim=1)
  53.                     gradient_penalty = 10 * gradient_norm.sub(1).pow(2).mean()
  54.                     loss_d += gradient_penalty
  55.  
  56.             # Discriminator update
  57.             self.loss += loss_d.item()
  58.             loss_d.backward()
  59.             self.optimizer.step()
  60.  
  61.             if self.gan_type == 'WGAN':
  62.                 for p in self.discriminator.parameters():
  63.                     p.data.clamp_(-1, 1)
  64.  
  65.         self.loss /= self.gan_k
  66.  
  67.         d_fake_for_g = self.discriminator(fake)
  68.         if self.gan_type == 'GAN':
  69.             loss_g = F.binary_cross_entropy_with_logits(
  70.                 d_fake_for_g, label_real
  71.             )
  72.         elif self.gan_type.find('WGAN') >= 0:
  73.             loss_g = -d_fake_for_g.mean()
  74.  
  75.         # Generator loss
  76.         return loss_g
  77.     
  78.     def state_dict(self, *args, **kwargs):
  79.         state_discriminator = self.discriminator.state_dict(*args, **kwargs)
  80.         state_optimizer = self.optimizer.state_dict()
  81.  
  82.         return dict(**state_discriminator, **state_optimizer)
  83.                
  84. # Some references
  85. # https://github.com/kuc2477/pytorch-wgan-gp/blob/master/model.py
  86. # OR
  87. # https://github.com/caogang/wgan-gp/blob/master/gan_cifar10.py

discriminator.py 

  1. from model import common
  2.  
  3. import torch.nn as nn
  4.  
  5. class Discriminator(nn.Module):
  6.     def __init__(self, args, gan_type='GAN'):
  7.         super(Discriminator, self).__init__()
  8.  
  9.         in_channels = 3
  10.         out_channels = 64
  11.         depth = 7
  12.         #bn = not gan_type == 'WGAN_GP'
  13.         bn = True
  14.         act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
  15.  
  16.         m_features = [
  17.             common.BasicBlock(args.n_colors, out_channels, 3, bn=bn, act=act)
  18.         ]
  19.         for i in range(depth):
  20.             in_channels = out_channels
  21.             if i % 2 == 1:
  22.                 stride = 1
  23.                 out_channels *= 2
  24.             else:
  25.                 stride = 2
  26.             m_features.append(common.BasicBlock(
  27.                 in_channels, out_channels, 3, stride=stride, bn=bn, act=act
  28.             ))
  29.  
  30.         self.features = nn.Sequential(*m_features)
  31.  
  32.         patch_size = args.patch_size // (2**((depth + 1) // 2))
  33.         m_classifier = [
  34.             nn.Linear(out_channels * patch_size**2, 1024),
  35.             act,
  36.             nn.Linear(1024, 1)
  37.         ]
  38.         self.classifier = nn.Sequential(*m_classifier)
  39.  
  40.     def forward(self, x):
  41.         features = self.features(x)
  42.         output = self.classifier(features.view(features.size(0), -1))
  43.  
  44.         return output
  45.

vgg.py 

  1. from model import common
  2.  
  3. import torch
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import torchvision.models as models
  7. from torch.autograd import Variable
  8.  
  9. class VGG(nn.Module):
  10.     def __init__(self, conv_index, rgb_range=1):
  11.         super(VGG, self).__init__()
  12.         # pretrained = True 表示使用已经训练过的参数
  13.         vgg_features = models.vgg19(pretrained=True).features
  14.         modules = [m for m in vgg_features]
  15.         if conv_index == '22':
  16.             self.vgg = nn.Sequential(*modules[:8])
  17.         elif conv_index == '54':
  18.             self.vgg = nn.Sequential(*modules[:35])
  19.  
  20.         vgg_mean = (0.485, 0.456, 0.406)
  21.         vgg_std = (0.229 * rgb_range, 0.224 * rgb_range, 0.225 * rgb_range)
  22.         self.sub_mean = common.MeanShift(rgb_range, vgg_mean, vgg_std)
  23.         self.vgg.requires_grad = False
  24.  
  25.     def forward(self, sr, hr):
  26.         def _forward(x):
  27.             x = self.sub_mean(x)
  28.             x = self.vgg(x)
  29.             return x
  30.             
  31.         vgg_sr = _forward(sr)
  32.         with torch.no_grad():
  33.             vgg_hr = _forward(hr.detach())
  34.  
  35.         loss = F.mse_loss(vgg_sr, vgg_hr)
  36.  
  37.         return loss

2.3.网络模型构建

dataloader.py

  1. import sys
  2. import threading
  3. import queue
  4. import random
  5. import collections
  6.  
  7. import torch
  8. import torch.multiprocessing as multiprocessing
  9.  
  10. from torch._C import _set_worker_signal_handlers, _update_worker_pids, \
  11.     _remove_worker_pids, _error_if_any_worker_fails
  12. from torch.utils.data.dataloader import DataLoader
  13. from torch.utils.data.dataloader import _DataLoaderIter
  14.  
  15. from torch.utils.data.dataloader import ExceptionWrapper
  16. from torch.utils.data.dataloader import _use_shared_memory
  17. from torch.utils.data.dataloader import _worker_manager_loop
  18. from torch.utils.data.dataloader import numpy_type_map
  19. from torch.utils.data.dataloader import default_collate
  20. from torch.utils.data.dataloader import pin_memory_batch
  21. from torch.utils.data.dataloader import _SIGCHLD_handler_set
  22. from torch.utils.data.dataloader import _set_SIGCHLD_handler
  23.  
  24. if sys.version_info[0] == 2:
  25.     import Queue as queue
  26. else:
  27.     import queue
  28.  
  29. def _ms_loop(dataset, index_queue, data_queue, collate_fn, scale, seed, init_fn, worker_id):
  30.     global _use_shared_memory
  31.     _use_shared_memory = True
  32.     _set_worker_signal_handlers()
  33.  
  34.     torch.set_num_threads(1)
  35.     torch.manual_seed(seed)
  36.     while True:
  37.         r = index_queue.get()
  38.         if r is None:
  39.             break
  40.         idx, batch_indices = r
  41.         try:
  42.             idx_scale = 0
  43.             if len(scale) > 1 and dataset.train:
  44.                 idx_scale = random.randrange(0, len(scale))
  45.                 dataset.set_scale(idx_scale)
  46.  
  47.             samples = collate_fn([dataset[i] for i in batch_indices])
  48.             samples.append(idx_scale)
  49.  
  50.         except Exception:
  51.             data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
  52.         else:
  53.             data_queue.put((idx, samples))
  54.  
  55. class _MSDataLoaderIter(_DataLoaderIter):
  56.     def __init__(self, loader):
  57.         self.dataset = loader.dataset
  58.         self.scale = loader.scale
  59.         self.collate_fn = loader.collate_fn
  60.         self.batch_sampler = loader.batch_sampler
  61.         self.num_workers = loader.num_workers
  62.         self.pin_memory = loader.pin_memory and torch.cuda.is_available()
  63.         self.timeout = loader.timeout
  64.         self.done_event = threading.Event()
  65.  
  66.         self.sample_iter = iter(self.batch_sampler)
  67.  
  68.         if self.num_workers > 0:
  69.             self.worker_init_fn = loader.worker_init_fn
  70.             self.index_queues = [
  71.                 multiprocessing.Queue() for _ in range(self.num_workers)
  72.             ]
  73.             self.worker_queue_idx = 0
  74.             self.worker_result_queue = multiprocessing.SimpleQueue()
  75.             self.batches_outstanding = 0
  76.             self.worker_pids_set = False
  77.             self.shutdown = False
  78.             self.send_idx = 0
  79.             self.rcvd_idx = 0
  80.             self.reorder_dict = {}
  81.  
  82.             base_seed = torch.LongTensor(1).random_()[0]
  83.             self.workers = [
  84.                 multiprocessing.Process(
  85.                     target=_ms_loop,
  86.                     args=(
  87.                         self.dataset,
  88.                         self.index_queues[i],
  89.                         self.worker_result_queue,
  90.                         self.collate_fn,
  91.                         self.scale,
  92.                         base_seed + i,
  93.                         self.worker_init_fn,
  94.                         i
  95.                     )
  96.                 )
  97.                 for i in range(self.num_workers)]
  98.  
  99.             if self.pin_memory or self.timeout > 0:
  100.                 self.data_queue = queue.Queue()
  101.                 if self.pin_memory:
  102.                     maybe_device_id = torch.cuda.current_device()
  103.                 else:
  104.                     # do not initialize cuda context if not necessary
  105.                     maybe_device_id = None
  106.                 self.worker_manager_thread = threading.Thread(
  107.                     target=_worker_manager_loop,
  108.                     args=(self.worker_result_queue, self.data_queue, self.done_event, self.pin_memory,
  109.                           maybe_device_id))
  110.                 self.worker_manager_thread.daemon = True
  111.                 self.worker_manager_thread.start()
  112.             else:
  113.                 self.data_queue = self.worker_result_queue
  114.  
  115.             for w in self.workers:
  116.                 w.daemon = True  # ensure that the worker exits on process exit
  117.                 w.start()
  118.  
  119.             _update_worker_pids(id(self), tuple(w.pid for w in self.workers))
  120.             _set_SIGCHLD_handler()
  121.             self.worker_pids_set = True
  122.  
  123.             # prime the prefetch loop
  124.             for _ in range(2 * self.num_workers):
  125.                 self._put_indices()
  126.  
  127. class MSDataLoader(DataLoader):
  128.     def __init__(
  129.         self, args, dataset, batch_size=1, shuffle=False,
  130.         sampler=None, batch_sampler=None,
  131.         collate_fn=default_collate, pin_memory=False, drop_last=False,
  132.         timeout=0, worker_init_fn=None):
  133.  
  134.         super(MSDataLoader, self).__init__(
  135.             dataset, batch_size=batch_size, shuffle=shuffle,
  136.             sampler=sampler, batch_sampler=batch_sampler,
  137.             num_workers=args.n_threads, collate_fn=collate_fn,
  138.             pin_memory=pin_memory, drop_last=drop_last,
  139.             timeout=timeout, worker_init_fn=worker_init_fn)
  140.  
  141.         self.scale = args.scale
  142.  
  143.     def __iter__(self):
  144.         return _MSDataLoaderIter(self)

 main.py

  1. import torch
  2.  
  3. import utility
  4. import data
  5. import model
  6. import loss
  7. from option import args
  8. from trainer import Trainer
  9.  
  10. torch.manual_seed(args.seed)
  11. checkpoint = utility.checkpoint(args)
  12.  
  13. if checkpoint.ok:
  14.     loader = data.Data(args)
  15.     model = model.Model(args, checkpoint)
  16.     loss = loss.Loss(args, checkpoint) if not args.test_only else None
  17.     t = Trainer(args, loader, model, loss, checkpoint)
  18.     while not t.terminate():
  19.         t.train()
  20.         t.test()
  21.  
  22.     checkpoint.done()
  23.

 option.py

  1. import argparse
  2. import template
  3.  
  4. parser = argparse.ArgumentParser(description='EDSR and MDSR')
  5.  
  6. parser.add_argument('--debug', action='store_true',
  7.                     help='Enables debug mode')
  8. parser.add_argument('--template', default='.',
  9.                     help='You can set various templates in option.py')
  10.  
  11. # Hardware specifications
  12. parser.add_argument('--n_threads', type=int, default=3,
  13.                     help='number of threads for data loading')
  14. parser.add_argument('--cpu', action='store_true',
  15.                     help='use cpu only')
  16. parser.add_argument('--n_GPUs', type=int, default=1,
  17.                     help='number of GPUs')
  18. parser.add_argument('--seed', type=int, default=1,
  19.                     help='random seed')
  20.  
  21. # Data specifications
  22. parser.add_argument('--dir_data', type=str, default='/home/yulun/data/SR/traindata/DIV2K/bicubic',
  23.                     help='dataset directory')
  24. parser.add_argument('--dir_demo', type=str, default='../test',
  25.                     help='demo image directory')
  26. parser.add_argument('--data_train', type=str, default='DIV2K',
  27.                     help='train dataset name')
  28. parser.add_argument('--data_test', type=str, default='DIV2K',
  29.                     help='test dataset name')
  30. parser.add_argument('--benchmark_noise', action='store_true',
  31.                     help='use noisy benchmark sets')
  32. parser.add_argument('--n_train', type=int, default=800,
  33.                     help='number of training set')
  34. parser.add_argument('--n_val', type=int, default=5,
  35.                     help='number of validation set')
  36. parser.add_argument('--offset_val', type=int, default=800,
  37.                     help='validation index offest')
  38. parser.add_argument('--ext', type=str, default='sep_reset',
  39.                     help='dataset file extension')
  40. parser.add_argument('--scale', default='4',
  41.                     help='super resolution scale')
  42. parser.add_argument('--patch_size', type=int, default=192,
  43.                     help='output patch size')
  44. parser.add_argument('--rgb_range', type=int, default=255,
  45.                     help='maximum value of RGB')
  46. parser.add_argument('--n_colors', type=int, default=3,
  47.                     help='number of color channels to use')
  48. parser.add_argument('--noise', type=str, default='.',
  49.                     help='Gaussian noise std.')
  50. parser.add_argument('--chop', action='store_true',
  51.                     help='enable memory-efficient forward')
  52.  
  53. # Model specifications
  54. parser.add_argument('--model', default='RCAN',
  55.                     help='model name')
  56.  
  57. parser.add_argument('--act', type=str, default='relu',
  58.                     help='activation function')
  59. parser.add_argument('--pre_train', type=str, default='.',
  60.                     help='pre-trained model directory')
  61. parser.add_argument('--extend', type=str, default='.',
  62.                     help='pre-trained model directory')
  63. parser.add_argument('--n_resblocks', type=int, default=20,
  64.                     help='number of residual blocks')
  65. parser.add_argument('--n_feats', type=int, default=64,
  66.                     help='number of feature maps')
  67. parser.add_argument('--res_scale', type=float, default=1,
  68.                     help='residual scaling')
  69. parser.add_argument('--shift_mean', default=True,
  70.                     help='subtract pixel mean from the input')
  71. parser.add_argument('--precision', type=str, default='single',
  72.                     choices=('single', 'half'),
  73.                     help='FP precision for test (single | half)')
  74.  
  75. # Training specifications
  76. parser.add_argument('--reset', action='store_true',
  77.                     help='reset the training')
  78. parser.add_argument('--test_every', type=int, default=1000,
  79.                     help='do test per every N batches')
  80. parser.add_argument('--epochs', type=int, default=1000,
  81.                     help='number of epochs to train')
  82. parser.add_argument('--batch_size', type=int, default=16,
  83.                     help='input batch size for training')
  84. parser.add_argument('--split_batch', type=int, default=1,
  85.                     help='split the batch into smaller chunks')
  86. parser.add_argument('--self_ensemble', action='store_true',
  87.                     help='use self-ensemble method for test')
  88. parser.add_argument('--test_only', action='store_true',
  89.                     help='set this option to test the model')
  90. parser.add_argument('--gan_k', type=int, default=1,
  91.                     help='k value for adversarial loss')
  92.  
  93. # Optimization specifications
  94. parser.add_argument('--lr', type=float, default=1e-4,
  95.                     help='learning rate')
  96. parser.add_argument('--lr_decay', type=int, default=200,
  97.                     help='learning rate decay per N epochs')
  98. parser.add_argument('--decay_type', type=str, default='step',
  99.                     help='learning rate decay type')
  100. parser.add_argument('--gamma', type=float, default=0.5,
  101.                     help='learning rate decay factor for step decay')
  102. parser.add_argument('--optimizer', default='ADAM',
  103.                     choices=('SGD', 'ADAM', 'RMSprop'),
  104.                     help='optimizer to use (SGD | ADAM | RMSprop)')
  105. parser.add_argument('--momentum', type=float, default=0.9,
  106.                     help='SGD momentum')
  107. parser.add_argument('--beta1', type=float, default=0.9,
  108.                     help='ADAM beta1')
  109. parser.add_argument('--beta2', type=float, default=0.999,
  110.                     help='ADAM beta2')
  111. parser.add_argument('--epsilon', type=float, default=1e-8,
  112.                     help='ADAM epsilon for numerical stability')
  113. parser.add_argument('--weight_decay', type=float, default=0,
  114.                     help='weight decay')
  115.  
  116. # Loss specifications
  117. parser.add_argument('--loss', type=str, default='1*L1',
  118.                     help='loss function configuration')
  119. parser.add_argument('--skip_threshold', type=float, default='1e6',
  120.                     help='skipping batch that has large error')
  121.  
  122. # Log specifications
  123. parser.add_argument('--save', type=str, default='test',
  124.                     help='file name to save')
  125. parser.add_argument('--load', type=str, default='.',
  126.                     help='file name to load')
  127. parser.add_argument('--resume', type=int, default=0,
  128.                     help='resume from specific checkpoint')
  129. parser.add_argument('--print_model', action='store_true',
  130.                     help='print model')
  131. parser.add_argument('--save_models', action='store_true',
  132.                     help='save all intermediate models')
  133. parser.add_argument('--print_every', type=int, default=100,
  134.                     help='how many batches to wait before logging training status')
  135. parser.add_argument('--save_results', action='store_true',
  136.                     help='save output results')
  137.  
  138. # options for residual group and feature channel reduction
  139. parser.add_argument('--n_resgroups', type=int, default=10,
  140.                     help='number of residual groups')
  141. parser.add_argument('--reduction', type=int, default=16,
  142.                     help='number of feature maps reduction')
  143. # options for test
  144. parser.add_argument('--testpath', type=str, default='../test/DIV2K_val_LR_our',
  145.                     help='dataset directory for testing')
  146. parser.add_argument('--testset', type=str, default='Set5',
  147.                     help='dataset name for testing')
  148.  
  149. args = parser.parse_args()
  150. template.set_template(args)
  151.  
  152. args.scale = list(map(lambda x: int(x), args.scale.split('+')))
  153.  
  154. if args.epochs == 0:
  155.     args.epochs = 1e8
  156.  
  157. for arg in vars(args):
  158.     if vars(args)[arg] == 'True':
  159.         vars(args)[arg] = True
  160.     elif vars(args)[arg] == 'False':
  161.         vars(args)[arg] = False
  162.

template.py 

  1. def set_template(args):
  2.     # Set the templates here
  3.     if args.template.find('jpeg') >= 0:
  4.         args.data_train = 'DIV2K_jpeg'
  5.         args.data_test = 'DIV2K_jpeg'
  6.         args.epochs = 200
  7.         args.lr_decay = 100
  8.  
  9.     if args.template.find('EDSR_paper') >= 0:
  10.         args.model = 'EDSR'
  11.         args.n_resblocks = 32
  12.         args.n_feats = 256
  13.         args.res_scale = 0.1
  14.  
  15.     if args.template.find('MDSR') >= 0:
  16.         args.model = 'MDSR'
  17.         args.patch_size = 48
  18.         args.epochs = 650
  19.  
  20.     if args.template.find('DDBPN') >= 0:
  21.         args.model = 'DDBPN'
  22.         args.patch_size = 128
  23.         args.scale = '4'
  24.  
  25.         args.data_test = 'Set5'
  26.  
  27.         args.batch_size = 20
  28.         args.epochs = 1000
  29.         args.lr_decay = 500
  30.         args.gamma = 0.1
  31.         args.weight_decay = 1e-4
  32.  
  33.         args.loss = '1*MSE'
  34.  
  35.     if args.template.find('GAN') >= 0:
  36.         args.epochs = 200
  37.         args.lr = 5e-5
  38.         args.lr_decay = 150
  39.

 utility.py

  1. import os
  2. import math
  3. import time
  4. import datetime
  5. from functools import reduce
  6.  
  7. import matplotlib
  8. matplotlib.use('Agg')
  9. import matplotlib.pyplot as plt
  10.  
  11. import numpy as np
  12. import scipy.misc as misc
  13.  
  14. import torch
  15. import torch.optim as optim
  16. import torch.optim.lr_scheduler as lrs
  17.  
  18. class timer():
  19.     def __init__(self):
  20.         self.acc = 0
  21.         self.tic()
  22.  
  23.     def tic(self):
  24.         self.t0 = time.time()
  25.  
  26.     def toc(self):
  27.         return time.time() - self.t0
  28.  
  29.     def hold(self):
  30.         self.acc += self.toc()
  31.  
  32.     def release(self):
  33.         ret = self.acc
  34.         self.acc = 0
  35.  
  36.         return ret
  37.  
  38.     def reset(self):
  39.         self.acc = 0
  40.  
  41. class checkpoint():
  42.     def __init__(self, args):
  43.         self.args = args
  44.         self.ok = True
  45.         self.log = torch.Tensor()
  46.         now = datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S')
  47.  
  48.         if args.load == '.':
  49.             if args.save == '.': args.save = now
  50.             self.dir = '../experiment/' + args.save
  51.         else:
  52.             self.dir = '../experiment/' + args.load
  53.             if not os.path.exists(self.dir):
  54.                 args.load = '.'
  55.             else:
  56.                 self.log = torch.load(self.dir + '/psnr_log.pt')
  57.                 print('Continue from epoch {}...'.format(len(self.log)))
  58.  
  59.         if args.reset:
  60.             os.system('rm -rf ' + self.dir)
  61.             args.load = '.'
  62.  
  63.         def _make_dir(path):
  64.             if not os.path.exists(path): os.makedirs(path)
  65.  
  66.         _make_dir(self.dir)
  67.         _make_dir(self.dir + '/model')
  68.         _make_dir(self.dir + '/results')
  69.  
  70.         open_type = 'a' if os.path.exists(self.dir + '/log.txt') else 'w'
  71.         self.log_file = open(self.dir + '/log.txt', open_type)
  72.         with open(self.dir + '/config.txt', open_type) as f:
  73.             f.write(now + '\n\n')
  74.             for arg in vars(args):
  75.                 f.write('{}: {}\n'.format(arg, getattr(args, arg)))
  76.             f.write('\n')
  77.  
  78.     def save(self, trainer, epoch, is_best=False):
  79.         trainer.model.save(self.dir, epoch, is_best=is_best)
  80.         trainer.loss.save(self.dir)
  81.         trainer.loss.plot_loss(self.dir, epoch)
  82.  
  83.         self.plot_psnr(epoch)
  84.         torch.save(self.log, os.path.join(self.dir, 'psnr_log.pt'))
  85.         torch.save(
  86.             trainer.optimizer.state_dict(),
  87.             os.path.join(self.dir, 'optimizer.pt')
  88.         )
  89.  
  90.     def add_log(self, log):
  91.         self.log = torch.cat([self.log, log])
  92.  
  93.     def write_log(self, log, refresh=False):
  94.         print(log)
  95.         self.log_file.write(log + '\n')
  96.         if refresh:
  97.             self.log_file.close()
  98.             self.log_file = open(self.dir + '/log.txt', 'a')
  99.  
  100.     def done(self):
  101.         self.log_file.close()
  102.  
  103.     def plot_psnr(self, epoch):
  104.         axis = np.linspace(1, epoch, epoch)
  105.         label = 'SR on {}'.format(self.args.data_test)
  106.         fig = plt.figure()
  107.         plt.title(label)
  108.         for idx_scale, scale in enumerate(self.args.scale):
  109.             plt.plot(
  110.                 axis,
  111.                 self.log[:, idx_scale].numpy(),
  112.                 label='Scale {}'.format(scale)
  113.             )
  114.         plt.legend()
  115.         plt.xlabel('Epochs')
  116.         plt.ylabel('PSNR')
  117.         plt.grid(True)
  118.         plt.savefig('{}/test_{}.pdf'.format(self.dir, self.args.data_test))
  119.         plt.close(fig)
  120.  
  121.     def save_results(self, filename, save_list, scale):
  122.         filename = '{}/results/{}_x{}_'.format(self.dir, filename, scale)
  123.         postfix = ('SR', 'LR', 'HR')
  124.         for v, p in zip(save_list, postfix):
  125.             normalized = v[0].data.mul(255 / self.args.rgb_range)
  126.             ndarr = normalized.byte().permute(1, 2, 0).cpu().numpy()
  127.             misc.imsave('{}{}.png'.format(filename, p), ndarr)
  128.  
  129. def quantize(img, rgb_range):
  130.     pixel_range = 255 / rgb_range
  131.     return img.mul(pixel_range).clamp(0, 255).round().div(pixel_range)
  132.  
  133. def calc_psnr(sr, hr, scale, rgb_range, benchmark=False):
  134.     diff = (sr - hr).data.div(rgb_range)
  135.     shave = scale
  136.     if diff.size(1) > 1:
  137.         convert = diff.new(1, 3, 1, 1)
  138.         convert[0, 0, 0, 0] = 65.738
  139.         convert[0, 1, 0, 0] = 129.057
  140.         convert[0, 2, 0, 0] = 25.064
  141.         diff.mul_(convert).div_(256)
  142.         diff = diff.sum(dim=1, keepdim=True)
  143.     '''
  144.     if benchmark:
  145.         shave = scale
  146.         if diff.size(1) > 1:
  147.             convert = diff.new(1, 3, 1, 1)
  148.             convert[0, 0, 0, 0] = 65.738
  149.             convert[0, 1, 0, 0] = 129.057
  150.             convert[0, 2, 0, 0] = 25.064
  151.             diff.mul_(convert).div_(256)
  152.             diff = diff.sum(dim=1, keepdim=True)
  153.     else:
  154.         shave = scale + 6
  155.     '''
  156.     valid = diff[:, :, shave:-shave, shave:-shave]
  157.     mse = valid.pow(2).mean()
  158.  
  159.     return -10 * math.log10(mse)
  160.  
  161. def make_optimizer(args, my_model):
  162.     trainable = filter(lambda x: x.requires_grad, my_model.parameters())
  163.  
  164.     if args.optimizer == 'SGD':
  165.         optimizer_function = optim.SGD
  166.         kwargs = {'momentum': args.momentum}
  167.     elif args.optimizer == 'ADAM':
  168.         optimizer_function = optim.Adam
  169.         kwargs = {
  170.             'betas': (args.beta1, args.beta2),
  171.             'eps': args.epsilon
  172.         }
  173.     elif args.optimizer == 'RMSprop':
  174.         optimizer_function = optim.RMSprop
  175.         kwargs = {'eps': args.epsilon}
  176.  
  177.     kwargs['lr'] = args.lr
  178.     kwargs['weight_decay'] = args.weight_decay
  179.     
  180.     return optimizer_function(trainable, **kwargs)
  181.  
  182. def make_scheduler(args, my_optimizer):
  183.     if args.decay_type == 'step':
  184.         scheduler = lrs.StepLR(
  185.             my_optimizer,
  186.             step_size=args.lr_decay,
  187.             gamma=args.gamma
  188.         )
  189.     elif args.decay_type.find('step') >= 0:
  190.         milestones = args.decay_type.split('_')
  191.         milestones.pop(0)
  192.         milestones = list(map(lambda x: int(x), milestones))
  193.         scheduler = lrs.MultiStepLR(
  194.             my_optimizer,
  195.             milestones=milestones,
  196.             gamma=args.gamma
  197.         )
  198.  
  199.     return scheduler
  200.

 trainer.py

  1. import os
  2. import math
  3. from decimal import Decimal
  4.  
  5. import utility
  6.  
  7. import torch
  8. from torch.autograd import Variable
  9. from tqdm import tqdm
  10.  
  11. class Trainer():
  12.     def __init__(self, args, loader, my_model, my_loss, ckp):
  13.         self.args = args
  14.         self.scale = args.scale
  15.  
  16.         self.ckp = ckp
  17.         self.loader_train = loader.loader_train
  18.         self.loader_test = loader.loader_test
  19.         self.model = my_model
  20.         self.loss = my_loss
  21.         self.optimizer = utility.make_optimizer(args, self.model)
  22.         self.scheduler = utility.make_scheduler(args, self.optimizer)
  23.  
  24.         if self.args.load != '.':
  25.             self.optimizer.load_state_dict(
  26.                 torch.load(os.path.join(ckp.dir, 'optimizer.pt'))
  27.             )
  28.             for _ in range(len(ckp.log)): self.scheduler.step()
  29.  
  30.         self.error_last = 1e8
  31.  
  32.     def train(self):
  33.         self.scheduler.step()
  34.         self.loss.step()
  35.         epoch = self.scheduler.last_epoch + 1
  36.         lr = self.scheduler.get_lr()[0]
  37.  
  38.         self.ckp.write_log(
  39.             '[Epoch {}]\tLearning rate: {:.2e}'.format(epoch, Decimal(lr))
  40.         )
  41.         self.loss.start_log()
  42.         self.model.train()
  43.  
  44.         timer_data, timer_model = utility.timer(), utility.timer()
  45.         for batch, (lr, hr, _, idx_scale) in enumerate(self.loader_train):
  46.             lr, hr = self.prepare([lr, hr])
  47.             timer_data.hold()
  48.             timer_model.tic()
  49.  
  50.             self.optimizer.zero_grad()
  51.             sr = self.model(lr, idx_scale)
  52.             loss = self.loss(sr, hr)
  53.             if loss.item() < self.args.skip_threshold * self.error_last:
  54.                 loss.backward()
  55.                 self.optimizer.step()
  56.             else:
  57.                 print('Skip this batch {}! (Loss: {})'.format(
  58.                     batch + 1, loss.item()
  59.                 ))
  60.  
  61.             timer_model.hold()
  62.  
  63.             if (batch + 1) % self.args.print_every == 0:
  64.                 self.ckp.write_log('[{}/{}]\t{}\t{:.1f}+{:.1f}s'.format(
  65.                     (batch + 1) * self.args.batch_size,
  66.                     len(self.loader_train.dataset),
  67.                     self.loss.display_loss(batch),
  68.                     timer_model.release(),
  69.                     timer_data.release()))
  70.  
  71.             timer_data.tic()
  72.  
  73.         self.loss.end_log(len(self.loader_train))
  74.         self.error_last = self.loss.log[-1, -1]
  75.  
  76.     def test(self):
  77.         epoch = self.scheduler.last_epoch + 1
  78.         self.ckp.write_log('\nEvaluation:')
  79.         self.ckp.add_log(torch.zeros(1, len(self.scale)))
  80.         self.model.eval()
  81.  
  82.         timer_test = utility.timer()
  83.         with torch.no_grad():
  84.             for idx_scale, scale in enumerate(self.scale):
  85.                 eval_acc = 0
  86.                 self.loader_test.dataset.set_scale(idx_scale)
  87.                 tqdm_test = tqdm(self.loader_test, ncols=80)
  88.                 for idx_img, (lr, hr, filename, _) in enumerate(tqdm_test):
  89.                     filename = filename[0]
  90.                     no_eval = (hr.nelement() == 1)
  91.                     if not no_eval:
  92.                         lr, hr = self.prepare([lr, hr])
  93.                     else:
  94.                         lr = self.prepare([lr])[0]
  95.  
  96.                     sr = self.model(lr, idx_scale)
  97.                     sr = utility.quantize(sr, self.args.rgb_range)
  98.  
  99.                     save_list = [sr]
  100.                     if not no_eval:
  101.                         eval_acc += utility.calc_psnr(
  102.                             sr, hr, scale, self.args.rgb_range,
  103.                             benchmark=self.loader_test.dataset.benchmark
  104.                         )
  105.                         save_list.extend([lr, hr])
  106.  
  107.                     if self.args.save_results:
  108.                         self.ckp.save_results(filename, save_list, scale)
  109.  
  110.                 self.ckp.log[-1, idx_scale] = eval_acc / len(self.loader_test)
  111.                 best = self.ckp.log.max(0)
  112.                 self.ckp.write_log(
  113.                     '[{} x{}]\tPSNR: {:.3f} (Best: {:.3f} @epoch {})'.format(
  114.                         self.args.data_test,
  115.                         scale,
  116.                         self.ckp.log[-1, idx_scale],
  117.                         best[0][idx_scale],
  118.                         best[1][idx_scale] + 1
  119.                     )
  120.                 )
  121.  
  122.         self.ckp.write_log(
  123.             'Total time: {:.2f}s\n'.format(timer_test.toc()), refresh=True
  124.         )
  125.         if not self.args.test_only:
  126.             self.ckp.save(self, epoch, is_best=(best[1][0] + 1 == epoch))
  127.  
  128.     def prepare(self, l, volatile=False):
  129.         device = torch.device('cpu' if self.args.cpu else 'cuda')
  130.         def _prepare(tensor):
  131.             if self.args.precision == 'half': tensor = tensor.half()
  132.             return tensor.to(device)
  133.            
  134.         return [_prepare(_l) for _l in l]
  135.  
  136.     def terminate(self):
  137.         if self.args.test_only:
  138.             self.test()
  139.             return True
  140.         else:
  141.             epoch = self.scheduler.last_epoch + 1
  142.             return epoch >= self.args.epochs
  143.

三.测试网络

利用模型将图片四倍放大的结果如下：

输入图片：

    

输出图片：

  

输入图片：

  

输出图片：

  

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/运维做开发/article/detail/782834
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号