使用SRGAN进行图像超分辨率

SRGAN（2倍）的生成器架构

子像素卷积中的像素重排

1import math
2import torch
3import torch.nn.functional as F
4from torch import nn
5import torchvision.models as models
6
7class ResidualBlock(nn.Module):
8    def __init__(self, channels):
9        super(ResidualBlock, self).__init__()
10        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
11        self.bn1 = nn.BatchNorm2d(channels)
12        self.prelu = nn.PReLU()
13        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
14        self.bn2 = nn.BatchNorm2d(channels)
15
16    def forward(self, x):
17        residual = self.conv1(x)
18        residual = self.bn1(residual)
19        residual = self.prelu(residual)
20        residual = self.conv2(residual)
21        residual = self.bn2(residual)
22        return x + residual

1class UpsampleBLock(nn.Module):
2    def __init__(self, in_channels, up_scale):
3        super(UpsampleBLock, self).__init__()
4        self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)
5        self.pixel_shuffle = nn.PixelShuffle(up_scale)
6        self.prelu = nn.PReLU()
7
8    def forward(self, x):
9        x = self.conv(x)
10        x = self.pixel_shuffle(x)
11        x = self.prelu(x)
12        return x

1class Generator(nn.Module):
2    def __init__(self, scale_factor):
3        upsample_block_num = int(math.log(scale_factor, 2))
4        super(Generator, self).__init__()
5        self.block1 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=9, padding=4), nn.PReLU())
6        self.block2 = ResidualBlock(64)
7        self.block3 = ResidualBlock(64)
8        self.block4 = ResidualBlock(64)
9        self.block5 = ResidualBlock(64)
10        self.block6 = ResidualBlock(64)
11        self.block7 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64))
12        block8 = [UpsampleBLock(64, 2) for _ in range(upsample_block_num)]
13        block8.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))
14        self.block8 = nn.Sequential(*block8)
15
16    def forward(self, x):
17        block1 = self.block1(x)
18        block2 = self.block2(block1)
19        block3 = self.block3(block2)
20        block4 = self.block4(block3)
21        block5 = self.block5(block4)
22        block6 = self.block6(block5)
23        block7 = self.block7(block6)
24        block8 = self.block8(block1 + block7)
25
26        return (torch.tanh(block8) + 1) / 2

SRGAN的判别器架构

1class Discriminator(nn.Module):
2    def __init__(self):
3        super(Discriminator, self).__init__()
4        self.net = nn.Sequential(
5            nn.Conv2d(3, 64, kernel_size=3, padding=1), nn.LeakyReLU(0.2),
6            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(64), nn.LeakyReLU(0.2),
7            nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2),
8            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2),
9            nn.Conv2d(128, 256, kernel_size=3, padding=1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2),
10            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2),
11            nn.Conv2d(256, 512, kernel_size=3, padding=1), nn.BatchNorm2d(512), nn.LeakyReLU(0.2),
12            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(512), nn.LeakyReLU(0.2),
13            nn.AdaptiveAvgPool2d(1), nn.Conv2d(512, 1024, kernel_size=1), nn.LeakyReLU(0.2), nn.Conv2d(1024, 1, kernel_size=1)
14        )
15
16    def forward(self, x):
17        batch_size = x.size(0)
18        return torch.sigmoid(self.net(x).view(batch_size))

$l_{\text{pixel}} + 10^{-3} \cdot l_{\text{adv}} + 6 \cdot 10^{-3} \cdot l_{\text{vgg}} + 2 \cdot 10^{-8} \cdot l_{\text{reg}}$

1class GeneratorLoss(nn.Module):
2    def __init__(self):
3        super(GeneratorLoss, self).__init__()
4        vgg = models.vgg16(pretrained=True)
5        loss_network = nn.Sequential(*list(vgg.features)[:31]).eval()
6        for param in loss_network.parameters():
7            param.requires_grad = False
8        self.loss_network = loss_network
9        self.mse_loss = nn.MSELoss()
10        self.tv_loss = TVLoss()
11
12    def forward(self, out_labels, out_images, target_images):
13        # 对抗损失
14        adversarial_loss = torch.mean(1 - out_labels)
15        # VGG 损失
16        vgg_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))
17        # 像素级损失
18        pixel_loss = self.mse_loss(out_images, target_images)
19        # 正则化损失
20        reg_loss = self.l2_loss(out_images)
21        return pixel_loss + 0.001 * adversarial_loss + 0.006 * vgg_loss + 2e-8 * reg_loss

1class L2Loss(nn.Module):
2    def __init__(self, l2_loss_weight=1):
3        super(L2Loss, self).__init__()
4        self.l2_loss_weight = l2_loss_weight
5
6    def forward(self, x):
7        batch_size = x.size()[0]
8        h_x = x.size()[2]
9        w_x = x.size()[3]
10        count_h = self.tensor_size(x[:, :, 1:, :])
11        count_w = self.tensor_size(x[:, :, :, 1:])
12        h_l2 = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
13        w_l2 = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()
14        return self.l2_loss_weight * 2 * (h_l2 / count_h + w_l2 / count_w) / batch_size
15
16    @staticmethod
17    def tensor_size(t):
18        return t.size()[1] * t.size()[2] * t.size()[3]

1# from loss import GeneratorLoss
2# from model import Generator, Discriminator
3from srgan1 import GeneratorLoss, Discriminator, Generator

1import argparse
2import os
3from math import log10
4
5import pandas as pd
6import torch.optim as optim
7import torch.utils.data
8import torchvision.utils as utils
9from torch.autograd import Variable
10from torch.utils.data import DataLoader
11from tqdm import tqdm
12
13import pytorch_ssim
14from data_utils import TrainDatasetFromFolder, ValDatasetFromFolder, display_transform
15from loss import GeneratorLoss
16from model import Generator, Discriminator
17
18parser = argparse.ArgumentParser(description='训练超分辨率模型')
19parser.add_argument('--crop_size', default=88, type=int, help='训练图像裁剪尺寸')
20parser.add_argument('--upscale_factor', default=4, type=int, choices=[2, 4, 8],
21                    help='超分辨率放大因子')
22parser.add_argument('--num_epochs', default=100, type=int, help='训练周期数')
23
24if __name__ == '__main__':
25    opt = parser.parse_args()
26    
27    CROP_SIZE = opt.crop_size
28    UPSCALE_FACTOR = opt.upscale_factor
29    NUM_EPOCHS = opt.num_epochs
30    
31    train_set = TrainDatasetFromFolder('data/DIV2K_train_HR', crop_size=CROP_SIZE, upscale_factor=UPSCALE_FACTOR)
32    val_set = ValDatasetFromFolder('data/DIV2K_valid_HR', upscale_factor=UPSCALE_FACTOR)
33    train_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=64, shuffle=True)
34    val_loader = DataLoader(dataset=val_set, num_workers=4, batch_size=1, shuffle=False)
35    
36    netG = Generator(UPSCALE_FACTOR)
37    print('# 生成器参数数量:', sum(param.numel() for param in netG.parameters()))
38    netD = Discriminator()
39    print('# 判别器参数数量:', sum(param.numel() for param in netD.parameters()))
40    
41    generator_criterion = GeneratorLoss()
42    
43    if torch.cuda.is_available():
44        netG.cuda()
45        netD.cuda()
46        generator_criterion.cuda()
47    
48    optimizerG = optim.Adam(netG.parameters())
49    optimizerD = optim.Adam(netD.parameters())
50    
51    results = {'d_loss': [], 'g_loss': [], 'd_score': [], 'g_score': [], 'psnr': [], 'ssim': []}
52    
53    for epoch in range(1, NUM_EPOCHS + 1):
54        train_bar = tqdm(train_loader)
55        running_results = {'batch_sizes': 0, 'd_loss': 0, 'g_loss': 0, 'd_score': 0, 'g_score': 0}
56    
57        netG.train()
58        netD.train()
59        for data, target in train_bar:
60            g_update_first = True
61            batch_size = data.size(0)
62            running_results['batch_sizes'] += batch_size
63    
64            ############################
65            # (1) 更新判别器网络：最大化 D(x) - D(G(z))
66            ###########################
67            real_img = Variable(target)
68            if torch.cuda.is_available():
69                real_img = real_img.cuda()
70            z = Variable(data)
71            if torch.cuda.is_available():
72                z = z.cuda()
73            fake_img = netG(z)
74    
75            netD.zero_grad()
76            real_out = netD(real_img).mean()
77            fake_out = netD(fake_img).mean()
78            d_loss = 1 - real_out + fake_out
79            d_loss.backward(retain_graph=True)
80            optimizerD.step()
81    
82            ############################
83            # (2) 更新生成器网络：最小化 1-D(G(z)) + 感知损失 + 图像损失 + TV损失
84            ###########################
85            netG.zero_grad()
86            ## 以下两行用于防止在Google Colab中出现运行时错误 ##
87            fake_img = netG(z)
88            fake_out = netD(fake_img).mean()
89            ##
90            g_loss = generator_criterion(fake_out, fake_img, real_img)
91            g_loss.backward()
92            
93            fake_img = netG(z)
94            fake_out = netD(fake_img).mean()
95            
96            optimizerG.step()
97
98            # 当前批次优化前的损失
99            running_results['g_loss'] += g_loss.item() * batch_size
100            running_results['d_loss'] += d_loss.item() * batch_size
101            running_results['d_score'] += real_out.item() * batch_size
102            running_results['g_score'] += fake_out.item() * batch_size
103    
104            train_bar.set_description(desc='[%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f' % (
105                epoch, NUM_EPOCHS, running_results['d_loss'] / running_results['batch_sizes'],
106                running_results['g_loss'] / running_results['batch_sizes'],
107                running_results['d_score'] / running_results['batch_sizes'],
108                running_results['g_score'] / running_results['batch_sizes']))
109    
110        netG.eval()
111        out_path = 'training_results/SRF_' + str(UPSCALE_FACTOR) + '/'
112        if not os.path.exists(out_path):
113            os.makedirs(out_path)
114        
115        with torch.no_grad():
116            val_bar = tqdm(val_loader)
117            valing_results = {'mse': 0, 'ssims': 0, 'psnr': 0, 'ssim': 0, 'batch_sizes': 0}
118            val_images = []
119            for val_lr, val_hr_restore, val_hr in val_bar:
120                batch_size = val_lr.size(0)
121                valing_results['batch_sizes'] += batch_size
122                lr = val_lr
123                hr = val_hr
124                if torch.cuda.is_available():
125                    lr = lr.cuda()
126                    hr = hr.cuda()
127                sr = netG(lr)
128        
129                batch_mse = ((sr - hr) ** 2).data.mean()
130                valing_results['mse'] += batch_mse * batch_size
131                batch_ssim = pytorch_ssim.ssim(sr, hr).item()
132                valing_results['ssims'] += batch_ssim * batch_size
133                valing_results['psnr'] = 10 * log10((hr.max()**2) / (valing_results['mse'] / valing_results['batch_sizes']))
134                valing_results['ssim'] = valing_results['ssims'] / valing_results['batch_sizes']
135                val_bar.set_description(
136                    desc='[将低分辨率图像转换为高分辨率图像] PSNR: %.4f dB SSIM: %.4f' % (
137                        valing_results['psnr'], valing_results['ssim']))
138        
139                val_images.extend(
140                    [display_transform()(val_hr_restore.squeeze(0)), display_transform()(hr.data.cpu().squeeze(0)),
141                     display_transform()(sr.data.cpu().squeeze(0))])
142            val_images = torch.stack(val_images)
143            val_images = torch.chunk(val_images, val_images.size(0) // 15)
144            val_save_bar = tqdm(val_images, desc='[保存训练结果]')
145            index = 1
146            for image in val_save_bar:
147                image = utils.make_grid(image, nrow=3, padding=5)
148                utils.save_image(image, out_path + 'epoch_%d_index_%d.png' % (epoch, index), padding=5)
149                index += 1
150    
151        # 保存模型参数
152        torch.save(netG.state_dict(), 'epochs/netG_epoch_%d_%d.pth' % (UPSCALE_FACTOR, epoch))
153        torch.save(netD.state_dict(), 'epochs/netD_epoch_%d_%d.pth' % (UPSCALE_FACTOR, epoch))
154        # 保存损失、得分、PSNR、SSIM
155        results['d_loss'].append(running_results['d_loss'] / running_results['batch_sizes'])
156        results['g_loss'].append(running_results['g_loss'] / running_results['batch_sizes'])
157        results['d_score'].append(running_results['d_score'] / running_results['batch_sizes'])
158        results['g_score'].append(running_results['g_score'] / running_results['batch_sizes'])
159        results['psnr'].append(valing_results['psnr'])
160        results['ssim'].append(valing_results['ssim'])
161    
162        if epoch % 10 == 0 and epoch != 0:
163            out_path = 'statistics/'
164            data_frame = pd.DataFrame(
165                data={'Loss_D': results['d_loss'], 'Loss_G': results['g_loss'], 'Score_D': results['d_score'],
166                      'Score_G': results['g_score'], 'PSNR': results['psnr'], 'SSIM': results['ssim']},
167                index=range(1, epoch + 1))
168            data_frame.to_csv(out_path + 'srf_' + str(UPSCALE_FACTOR) + '_train_results.csv', index_label='Epoch')