Add CI for Super Resolution example and tqdm bar to the example

.github/workflows/unit-tests.yml

@@ -184,3 +184,8 @@ jobs:
           #train
           mkdir -p ~/.cache/torch/checkpoints/ && wget "https://download.pytorch.org/models/vgg16-397923af.pth" -O ~/.cache/torch/checkpoints/vgg16-397923af.pth
           python examples/fast_neural_style/neural_style.py train --epochs 1 --cuda 0 --dataset test --dataroot . --image_size 32 --style_image examples/fast_neural_style/images/style_images/mosaic.jpg --style_size 32
+      - name: Run SR Example
+        if: ${{ matrix.os == 'ubuntu-latest' }}
+        run: |
+          # Super-Resolution
+          python examples/super_resolution/main.py --upscale_factor 3 --crop_size 180 --batch_size 4 --test_batch_size 100 --n_epochs 1 --lr 0.001 --threads 2

examples/super_resolution/README.md

@@ -5,7 +5,7 @@ ported from [pytorch-examples](https://github.com/pytorch/examples/tree/main/sup
 This example illustrates how to use the efficient sub-pixel convolution layer described in ["Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network" - Shi et al.](https://arxiv.org/abs/1609.05158) for increasing spatial resolution within your network for tasks such as superresolution.
 
 ```
-usage: main.py [-h] --upscale_factor UPSCALE_FACTOR [--batch_size BATCHSIZE]
+usage: main.py [-h] --upscale_factor UPSCALE_FACTOR [--crop_size CROPSIZE] [--batch_size BATCHSIZE]
                [--test_batch_size TESTBATCHSIZE] [--n_epochs NEPOCHS] [--lr LR]
                [--cuda] [--threads THREADS] [--seed SEED]
 
@@ -14,6 +14,7 @@ PyTorch Super Res Example
 optional arguments:
   -h, --help            show this help message and exit
   --upscale_factor      super resolution upscale factor
+  --crop_size           cropped size of the images for training
   --batch_size           training batch size
   --test_batch_size       testing batch size
   --n_epochs             number of epochs to train for
@@ -30,7 +31,7 @@ This example trains a super-resolution network on the [Caltech101 dataset](https
 
 ### Train
 
-`python main.py --upscale_factor 3 --batch_size 4 --test_batch_size 100 --n_epochs 30 --lr 0.001`
+`python main.py --upscale_factor 3 --crop_size 180 --batch_size 4 --test_batch_size 100 --n_epochs 30 --lr 0.001`
 
 ### Super Resolve
 

examples/super_resolution/main.py

@@ -8,11 +8,14 @@
 from torch.utils.data import DataLoader
 from torchvision.transforms.functional import center_crop, resize, to_tensor
 
+from ignite.contrib.handlers import ProgressBar
+
 from ignite.engine import Engine, Events
 from ignite.metrics import PSNR
 
 # Training settings
 parser = argparse.ArgumentParser(description="PyTorch Super Res Example")
+parser.add_argument("--crop_size", type=int, default=256, help="cropped size of the images for training")
 parser.add_argument("--upscale_factor", type=int, required=True, help="super resolution upscale factor")
 parser.add_argument("--batch_size", type=int, default=64, help="training batch size")
 parser.add_argument("--test_batch_size", type=int, default=10, help="testing batch size")
@@ -22,6 +25,7 @@
 parser.add_argument("--mps", action="store_true", default=False, help="enables macOS GPU training")
 parser.add_argument("--threads", type=int, default=4, help="number of threads for data loader to use")
 parser.add_argument("--seed", type=int, default=123, help="random seed to use. Default=123")
+
 opt = parser.parse_args()
 
 print(opt)
@@ -70,8 +74,8 @@ def __len__(self):
 trainset = torchvision.datasets.Caltech101(root="./data", download=True)
 testset = torchvision.datasets.Caltech101(root="./data", download=False)
 
-trainset_sr = SRDataset(trainset, scale_factor=opt.upscale_factor)
-testset_sr = SRDataset(testset, scale_factor=opt.upscale_factor)
+trainset_sr = SRDataset(trainset, scale_factor=opt.upscale_factor, crop_size=opt.crop_size)
+testset_sr = SRDataset(testset, scale_factor=opt.upscale_factor, crop_size=opt.crop_size)
 
 training_data_loader = DataLoader(dataset=trainset_sr, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True)
 testing_data_loader = DataLoader(dataset=testset_sr, num_workers=opt.threads, batch_size=opt.test_batch_size)
@@ -145,4 +149,6 @@ def checkpoint():
     print("Checkpoint saved to {}".format(model_out_path))
 
 
+ProgressBar().attach(trainer)
+
 trainer.run(training_data_loader, opt.n_epochs)