Shulou(Shulou.com)05/31 Report--

This article introduces the knowledge of "how to build Pytorch SRGAN platform to improve picture super-resolution". In the operation of actual cases, many people will encounter such a dilemma, so let the editor lead you to learn how to deal with these situations. I hope you can read it carefully and be able to achieve something!

Network Construction 1. What is SRGAN

SRGAN comes from the paper Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network.

If SRGAN is regarded as a black box, its main function is to input a low-resolution picture to generate a high-resolution picture.

The article mentioned that the ordinary super-resolution model only uses the mean square error as the loss function when training the network, although it can obtain a high peak signal-to-noise ratio, but the restored image usually loses high-frequency details.

SRGAN uses perceived loss (perceptual loss) and countermeasure loss (adversarial loss) to enhance the realism of the recovered image.

Second, the construction of the generation network

The composition of the generating network is shown above, and the function of the generating network is to input a low-resolution picture to generate a high-resolution picture. :

The generation network of SRGAN consists of three parts.

1. The low-resolution image will go through a convolution + RELU function.

2. Then through B residual network structure, each residual structure contains two convolution + standardization + RELU, and there is a residual edge.

3. Then enter the upsampling part, after two upsampling, the height and width of the original image becomes 4 times of the original, and the resolution is improved.

The first two parts are used for feature extraction, and the third part is used to improve the resolution.

Import mathimport torchfrom torch import nnclass ResidualBlock (nn.Module): def _ init__ (self, channels): super (ResidualBlock, self). _ _ init__ () self.conv1 = nn.Conv2d (channels, channels, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d (channels) self.prelu = nn.PReLU (channels) self.conv2 = nn.Conv2d (channels, channels, kernel_size=3 Padding=1) self.bn2 = nn.BatchNorm2d (channels) def forward (self, x): short_cut = x = self.conv1 (x) x = self.bn1 (x) x = self.prelu (x) x = self.conv2 (x) x = self.bn2 (x) return x + short_cutclass UpsampleBLock (nn.Module): def _ init__ (self, in_channels) Up_scale): super (UpsampleBLock, self). _ _ init__ () self.conv = nn.Conv2d (in_channels, in_channels * up_scale * * 2, kernel_size=3, padding=1) self.pixel_shuffle = nn.PixelShuffle (up_scale) self.prelu = nn.PReLU (in_channels) def forward (self X): X = self.conv (x) x = self.pixel_shuffle (x) x = self.prelu (x) return xclass Generator (nn.Module): def _ init__ (self, scale_factor, num_residual=16): upsample_block_num = int (math.log (scale_factor, 2)) super (Generator Self). _ init__ () self.block_in = nn.Sequential (nn.Conv2d (3,64, kernel_size=9, padding=4) Nn.PReLU (64)) self.blocks = [] for _ in range (num_residual): self.blocks.append (ResidualBlock (64)) self.blocks = nn.Sequential (* self.blocks) self.block_out = nn.Sequential (nn.Conv2d (64, 64, kernel_size=3, padding=1) Nn.BatchNorm2d (64)) self.upsample = [UpsampleBLock (64,2) for _ in range (upsample_block_num)] self.upsample.append (nn.Conv2d (64,3, kernel_size=9, padding=4)) self.upsample = nn.Sequential (* self.upsample) def forward (self) X): X = self.block_in (x) short_cut = x x = self.blocks (x) x = self.block_out (x) upsample = self.upsample (x + short_cut) return torch.tanh (upsample) 3. Construction of Discriminant Network

The composition of the discriminant network is shown in the figure above:

The discriminant network of SRGAN consists of repeated convolution + LeakyRELU and standardization.

For judging the network, its purpose is to judge whether the input picture is true or false, its input is the picture, and the output is the judgment result.

The result of the judgment is between 0 and 1, using close to 1 to judge as a true picture, and close to 0 as a false picture.

It is judged that there is little difference between the construction of the network and the ordinary convolution network, both of which are used under the convolution of the picture. After many convolution, the result of the judgment is fully connected one after another.

The implementation code is as follows:

Class Discriminator (nn.Module): def _ init__ (self): super (Discriminator, self). _ _ init__ () self.net = nn.Sequential (nn.Conv2d (3,64, kernel_size=3, padding=1), nn.LeakyReLU (0.2), nn.Conv2d (64, 64, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d (64) Nn.LeakyReLU, nn.Conv2d (64,128, kernel_size=3, padding=1), nn.BatchNorm2d (128), nn.LeakyReLU (0.2), nn.Conv2d (128,128, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d (128), nn.LeakyReLU (0.2), nn.Conv2d (128,256) Kernel_size=3, padding=1, nn.BatchNorm2d, nn.LeakyReLU, nn.Conv2d (256,256, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d, nn.LeakyReLU, nn.Conv2d (256,512, kernel_size=3, padding=1), nn.BatchNorm2d Nn.LeakyReLU, nn.Conv2d (512,512, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d, nn.LeakyReLU (0.2), nn.AdaptiveAvgPool2d (1), nn.Conv2d (512, 1024, kernel_size=1), nn.LeakyReLU (0.2), nn.Conv2d (1024, 1) Kernel_size=1)) def forward (self, x): batch_size = x.size (0) return torch.sigmoid (self.net (x) .view (batch_size)) training ideas

SRGAN training can be divided into generator training and discriminator training:

In each step, the discriminator is generally trained first, and then the generator is trained.

I. training of discriminators

When training the discriminator, we hope that the discriminator can determine whether the input image is true or false, so our input is the real picture, the fake picture and their corresponding tags.

Therefore, the training steps of the discriminator are as follows:

1. Randomly select batch_size real high-resolution images.

2. Use the low-resolution images after resize and transfer them into Generator to generate batch_size fake high-resolution images.

3. The label of the real picture is 1, and the label of the false picture is 0. The real picture and the false picture are introduced into the Discriminator as a training set for training.

Second, the training of generators

When training the generator, we hope that the generator can generate very real fake images. So when we train the generator, we need to know what the discriminator thinks is the real picture.

Therefore, the training steps of the generator are as follows:

1. The low-resolution image is introduced into the generation model to get the false high-resolution image, and the discrimination result of the false high-resolution image is compared with 1 to get loss. (compared with 1 means that the generator is trained according to the result of the discriminator.)

2. The real high-resolution image and the false high-resolution image are introduced into the VGG network to obtain the features of the two images, and the loss is obtained by comparing the features of the two images.

Use SRGAN to generate pictures

The overall library structure of SRGAN is as follows:

I. preparation of data sets

The dataset needs to be prepared before training, and the dataset is saved in the datasets folder.

II. Processing of data sets

Open txt_annotation.py, which points to the datasets in the root directory by default. Run txt_annotation.py.

The train_lines.txt under the root directory is generated at this point.

III. Model training

After completing the dataset processing, run train.py to start the training.

During training, you can view the training results in the results folder:

This is the end of the content of "how to build a SRGAN platform for Pytorch to improve picture super resolution". Thank you for reading. If you want to know more about the industry, you can follow the website, the editor will output more high-quality practical articles for you!

Welcome to subscribe "Shulou Technology Information " to get latest news, interesting things and hot topics in the IT industry, and controls the hottest and latest Internet news, technology news and IT industry trends.