Add the example of super_resolution

examples/super_resolution/README.md

@@ -30,7 +30,7 @@ This example trains a super-resolution network on the [Cifar10 dataset](https://
 
 ### Train
 
-`python main.py --upscale_factor 3 --batchSize 4 --testBatchSize 100 --nEpochs 30 --lr 0.001`
+`python main.py --upscale_factor 3 --batch_size 4 --test_batch_size 100 --n_epochs 30 --lr 0.001`
 
 ### Super Resolve
 

examples/super_resolution/main.py

@@ -4,19 +4,19 @@
 import torch.nn as nn
 import torch.optim as optim
 import torchvision
-import torchvision.transforms as transforms
 from model import Net
 from torch.utils.data import DataLoader
+from torchvision.transforms import CenterCrop, Compose, Resize, ToTensor
 
 from ignite.engine import Engine, Events
 from ignite.metrics import PSNR
 
 # Training settings
 parser = argparse.ArgumentParser(description="PyTorch Super Res Example")
 parser.add_argument("--upscale_factor", type=int, required=True, help="super resolution upscale factor")
-parser.add_argument("--batchSize", type=int, default=64, help="training batch size")
-parser.add_argument("--testBatchSize", type=int, default=10, help="testing batch size")
-parser.add_argument("--nEpochs", type=int, default=2, help="number of epochs to train for")
+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")
+parser.add_argument("--n_epochs", type=int, default=2, help="number of epochs to train for")
 parser.add_argument("--lr", type=float, default=0.01, help="Learning Rate. Default=0.01")
 parser.add_argument("--cuda", action="store_true", help="use cuda?")
 parser.add_argument("--mps", action="store_true", default=False, help="enables macOS GPU training")
@@ -45,32 +45,71 @@
 
 
 class SRDataset(torch.utils.data.Dataset):
-    def __init__(self, dataset, scale_factor):
+    def __init__(self, dataset, scale_factor, input_transform=None, target_transform=None):
         self.dataset = dataset
-        self.transform = transforms.Resize(
-            (len(dataset[0][0][0]) * scale_factor, len(dataset[0][0][0][0]) * scale_factor)
-        )
+        self.input_transform = input_transform
+        self.target_transform = target_transform
 
     def __getitem__(self, index):
-        lr_image, _ = self.dataset[index]
-        hr_image = self.transform(lr_image)
+        image, _ = self.dataset[index]
+        img = image.convert("YCbCr")
+        lr_image, _, _ = img.split()
+
+        hr_image = lr_image.copy()
+        if self.input_transform:
+            lr_image = self.input_transform(lr_image)
+        if self.target_transform:
+            hr_image = self.target_transform(hr_image)
         return lr_image, hr_image
 
     def __len__(self):
         return len(self.dataset)
 
 
-transform = transforms.Compose([transforms.ToTensor()])
+def calculate_valid_crop_size(crop_size, upscale_factor):
+    return crop_size - (crop_size % upscale_factor)
+
+
+def input_transform(crop_size, upscale_factor):
+    return Compose(
+        [
+            CenterCrop(crop_size),
+            Resize(crop_size // upscale_factor),
+            ToTensor(),
+        ]
+    )
+
+
+def target_transform(crop_size):
+    return Compose(
+        [
+            CenterCrop(crop_size),
+            ToTensor(),
+        ]
+    )
+
+
+crop_size = calculate_valid_crop_size(256, opt.upscale_factor)
 
-trainset = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
-testset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True, transform=transform)
+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,
+    input_transform=input_transform(crop_size, opt.upscale_factor),
+    target_transform=target_transform(crop_size),
+)
+testset_sr = SRDataset(
+    testset,
+    scale_factor=opt.upscale_factor,
+    input_transform=input_transform(crop_size, opt.upscale_factor),
+    target_transform=target_transform(crop_size),
+)
 
-training_data_loader = DataLoader(dataset=trainset_sr, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
+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.testBatchSize, shuffle=False
+    dataset=testset_sr, num_workers=opt.threads, batch_size=opt.test_batch_size, shuffle=False
 )
 
 print("===> Building model")
@@ -105,7 +144,7 @@ def validation_step(engine, batch):
 psnr = PSNR(data_range=1)
 psnr.attach(evaluator, "psnr")
 validate_every = 1
-log_interval = 10
+log_interval = 100
 
 
 @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
@@ -117,13 +156,9 @@ def log_training_loss(engine):
     )
 
 
-@trainer.on(Events.EPOCH_COMPLETED(every=validate_every))
-def run_validation():
-    evaluator.run(testing_data_loader)
-
-
 @trainer.on(Events.EPOCH_COMPLETED(every=validate_every))
 def log_validation():
+    evaluator.run(testing_data_loader)
     metrics = evaluator.state.metrics
     print(f"Epoch: {trainer.state.epoch}, Avg. PSNR: {metrics['psnr']} dB")
 
@@ -145,4 +180,4 @@ def checkpoint():
     print("Checkpoint saved to {}".format(model_out_path))
 
 
-trainer.run(training_data_loader, opt.nEpochs)
+trainer.run(training_data_loader, opt.n_epochs)