PyTorch Tutorial to Super-Resolution

This repository provides an updated implementation of Super-Resolution (SR) using Generative Adversarial Networks (GANs) in PyTorch 1.4 with Python 3.6, inspired by the original work in Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network. It implements SRResNet and SRGAN to upscale low-resolution images by 4x (16x pixel count) for photorealistic results. This project builds upon and modernizes the foundational work by sgrvinod/a-PyTorch-Tutorial-to-Super-Resolution, to which we extend our gratitude for its comprehensive inspiration.

Assumptions: Familiarity with PyTorch and convolutional neural networks (CNNs). For PyTorch beginners, refer to Deep Learning with PyTorch: A 60 Minute Blitz.

Feedback: Please submit questions or issues on the GitHub Issues page.

Contents

• Objective
• Concepts
• Implementation
• Training
• Evaluation
• Inference
• FAQs

Objective

The objective is to develop models that upscale low-resolution images by 4x, generating photorealistic high-resolution images using:

• SRResNet: A deep CNN leveraging residual connections.
• SRGAN: Enhances SRResNet through adversarial training for improved realism.

View example outputs (e.g., earth, baboon, flowers) in the original repository by sgrvinod or the SRGAN video demo in 1080p.

Concepts

• Super-Resolution: Generates new pixels to enhance image resolution.
• Residual Connections: Skip connections improve deep network training (He et al., 2015).
• GANs: A Generator produces realistic images, while a Discriminator differentiates real from fake (Goodfellow et al., 2014).
• Sub-Pixel Convolution: Upscales by generating channels, then shuffling into a higher-resolution image.
• Perceptual Loss: Combines MSE-based content loss in VGG feature space with adversarial loss for photorealism.

Implementation

Dataset

• Training: ~120k MSCOCO 2017 images (train2017, val2017).
• Testing: Set5, Set14, BSD100 datasets.
• Download:
  • MSCOCO: Training (19GB), Validation (1GB).
  • Test datasets: Set5, Set14, BSD100 (use original folder in Set5/Set14).
• Organize in folders: train2017, val2017, Set5, Set14, BSD100.

Model Inputs and Targets

• High-Resolution (HR) Images: 96x96 patches, normalized to [-1, 1] (SRResNet) or ImageNet-normed (SRGAN).
• Low-Resolution (LR) Images: 4x downsampled HR images (24x24), ImageNet-normed.
• Super-Resolved (SR) Images: SRResNet/Generator outputs, 96x96, in [-1, 1].
• Discriminator Labels: Binary (1 for HR, 0 for SR).

Data Pipeline

• Parsing: create_data_lists.py generates JSON files with image filepaths.
• Transforms: HR images are random 96x96 crops (training) or center-crops (testing). LR images are 4x bicubically downsampled with antialiasing (Pillow).
• Dataset: SRDataset (datasets.py) returns LR/HR pairs.
• DataLoader: Feeds batches for training/evaluation.

Models

• ConvolutionalBlock (models.py): 2D convolution, optional batch norm, Tanh/PReLU/LeakyReLU.
• SubPixelConvolutionalBlock (models.py): Convolution to s²n channels + nn.PixelShuffle for upscaling.
• ResidualBlock (models.py): Two ConvolutionalBlocks with a skip connection.
• SRResNet (models.py): 16 residual blocks, 2 sub-pixel convolution blocks for 4x upscaling, outputs [-1, 1].
• Generator (models.py): Identical to SRResNet, initialized with trained SRResNet weights.
• Discriminator (models.py): Binary classifier with 7 convolutional blocks, doubles channels in even blocks, halves resolution in odd blocks, outputs probability.
• TruncatedVGG19 (models.py): Pretrained VGG19 truncated at 5th maxpool, 4th convolution for content loss.

Training

1. Prepare Data: Run create_data_lists.py.

   python create_data_lists.py

2. Train SRResNet:
   • Script: train_srresnet.py.
   • Hyperparameters: Adam optimizer, lr=1e-4, 1M iterations, batch size=16.
   • Run: python train_srresnet.py (resume with checkpoint path).
3. Train SRGAN:
   • Script: train_srgan.py.
   • Requires SRResNet checkpoint for Generator initialization.
   • Hyperparameters: Adam optimizer, lr=1e-4 for 100k iterations, lr=1e-5 for 100k more, batch size=16.
   • Run: python train_srgan.py (resume with checkpoint path).
4. Checkpoints: Download pretrained models here.

Evaluation

• Script: eval.py.
• Metrics: PSNR, SSIM on Set5, Set14, BSD100.
• Run: python eval.py.
• Results:

  Model	Set5 PSNR	Set5 SSIM	Set14 PSNR	Set14 SSIM	BSD100 PSNR	BSD100 SSIM
  SRResNet	31.927	0.902	28.588	0.799	27.587	0.756
  SRGAN	29.719	0.859	26.509	0.729	25.531	0.678

• Note: PSNR/SSIM favor SRResNet’s smoother images; SRGAN excels in photorealism per user studies.

Inference

• Script: super_resolve.py.
• Function: visualize_sr() compares HR, bicubic-upsampled LR, SRResNet, and SRGAN outputs.
• Run: Specify HR image and checkpoints in super_resolve.py.
• Note: Use halve=True for large images to fit screen (e.g., 2160p HR → 1080p HR/SR, 540p LR → 270p LR on 1080p screen).

Frequently Asked Questions

1. Why pass SR images through Discriminator twice?
   • First pass (Generator training): SR images in computational graph for adversarial loss, backpropagating through Generator.
   • Second pass (Discriminator training): SR images detached to avoid Generator backpropagation, saving computation.
2. Subpixel vs. Transposed Convolution?
   • Subpixel convolution creates s²n channels, shuffles into upscaled image, often more efficient. May optimize better (see discussion).

Acknowledgments

This repository is an update and extension of the foundational work by sgrvinod/a-PyTorch-Tutorial-to-Super-Resolution. We express our appreciation for the original implementation, which served as a critical inspiration for this project.

For full code and examples, visit this repository or the original sgrvinod/a-PyTorch-Tutorial-to-Super-Resolution.