Update neatly organized codes

2KangHo · 2KangHo · commit e27e933ca158 · 2019-02-14T14:30:33.000+09:00
diff --git a/README.md b/README.md
@@ -1,41 +1,30 @@
-# vdsr_pytorch
-
-PyTorch Implementation  
-no muliti-scale  
-but you can use multi-GPUs
+# vdsr_pytorch_lms
+VDSR PyTorch Implementation  
+You can use multi-gpus.  
+but no multi-scale.
+And you can input gaussian noise to input images.
 
 ## Requirement
-
 `torch`  
 `torchvision`  
-`python-tk` (or `pyhton3-tk`)
-
-## Training
+`python-tk` (or `python3-tk`)
 
-    > python main.py --batch_size 40 --test_batch_size 10 --epochs 100 --cuda --gpuids 0 --upscale_factor 2
+## Download dataset
+1. Download [DF2K dataset](https://drive.google.com/file/d/1P9pcaGjvq3xiF22GXIq7ciZta3rjZxaY/view?usp=sharing).
+2. move under dataset directory i.e. vdsr_pytorch_lms/dataset/DF2K
 
-or
-
-    > python3 main.py --batch_size 40 --test_batch_size 10 --epochs 100 --cuda --gpuids 0 --upscale_factor 2
+## Training
+```
+$ python main.py --dataset DF2K --cuda --gpuids 0 1 --upscale_factor 2 --crop_size 256 --batch_size 128 --test_batch_size 32
+```
 
 ## Test
+```
+$ python main.py --dataset Urban100 --cuda --gpuids 0 1 --upscale_factor 2 --crop_size 256 --test_batch_size 32 --test --model model_epoch_100.pth
+```
 
-    > python main.py --batch_size 40 --test_batch_size 10 --epochs 100 --cuda --gpuids 0 --upscale_factor 2 --test --model model_epoch_100.pth
-
-or
-
-    > python3 main.py --batch_size 40 --test_batch_size 10 --epochs 100 --cuda --gpuids 0 --upscale_factor 2 --test --model model_epoch_100.pth
-
-## Sample Usage
-
-    > python run.py --input_image test_scale2x.jpg --scale_factor 2 --model model_epoch_100.pth --cuda --gpuids 0 --output_filename test_scale2x_out.jpg
-
-or
-
-    > python3 run.py --input_image test_scale2x.jpg --scale_factor 2 --model model_epoch_100.pth --cuda --gpuids 0 --output_filename test_scale2x_out.jpg
-
-## 주의
+## Sample usage
+```
+$ python run.py --cuda --gpuids 0 1 --scale_factor 2 --model model_epoch_100.pth --input_image test_scale2x.jpg --output_filename test_scale2x_out.jpg
+```
 
-test시에는 Urban100의 데이터가 사용된다.  
-sample에서 input image는 학습에 사용된 BSDS300 data가 아닌 인터넷에서 가져온 이미지 등 다른 이미지를 사용해야 정확한 성능을 확인할 수 있다.  
-Urban100데이터는 학습시 사용하지 않으므로 Urban100의 사진들 중 하나를 사용해도 됨.
diff --git a/data.py b/data.py
@@ -1,44 +1,17 @@
-from os.path import exists, join, basename
-from os import makedirs, remove
-from six.moves import urllib
-import tarfile
+from os.path import join
 from torchvision.transforms import Compose, CenterCrop, ToTensor, Resize
 
 from data_utils import DatasetFromFolder
 
 
-def download_bsd300(dest="dataset"):
-    output_image_dir = join(dest, "BSDS300/images")
-
-    if not exists(output_image_dir):
-        makedirs(dest)
-        url = "http://www2.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/BSDS300-images.tgz"
-        print("downloading url ", url)
-
-        data = urllib.request.urlopen(url)
-
-        file_path = join(dest, basename(url))
-        with open(file_path, 'wb') as f:
-            f.write(data.read())
-
-        print("Extracting data")
-        with tarfile.open(file_path) as tar:
-            for item in tar:
-                tar.extract(item, dest)
-
-        remove(file_path)
-
-    return output_image_dir
-
-
 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*2), crop_size//upscale_factor)),
+        Resize((crop_size//upscale_factor, crop_size//upscale_factor)),
         Resize((crop_size, crop_size)),
         ToTensor(),
     ])
@@ -51,36 +24,35 @@ def target_transform(crop_size):
     ])
 
 
-def get_training_set(upscale_factor, add_noise=None, noise_std=3.0):
-    root_dir = download_bsd300()
-
+def get_training_set(dataset, crop_size, upscale_factor, add_noise=None, noise_std=3.0):
+    root_dir = join("dataset", dataset)
     train_dir = join(root_dir, "train")
-    crop_size = calculate_valid_crop_size(256, upscale_factor)
+    cropsize = calculate_valid_crop_size(crop_size, upscale_factor)
 
     return DatasetFromFolder(train_dir,
                              input_transform=input_transform(
-                                 crop_size, upscale_factor),
-                             target_transform=target_transform(crop_size),
+                                 cropsize, upscale_factor),
+                             target_transform=target_transform(cropsize),
                              add_noise=add_noise,
                              noise_std=noise_std)
 
 
-def get_validation_set(upscale_factor):
-    root_dir = download_bsd300()
-    validation_dir = join(root_dir, "test")
-    crop_size = calculate_valid_crop_size(256, upscale_factor)
+def get_validation_set(dataset, crop_size, upscale_factor):
+    root_dir = join("dataset", dataset)
+    validation_dir = join(root_dir, "valid")
+    cropsize = calculate_valid_crop_size(crop_size, upscale_factor)
 
     return DatasetFromFolder(validation_dir,
                              input_transform=input_transform(
-                                 crop_size, upscale_factor),
-                             target_transform=target_transform(crop_size))
+                                 cropsize, upscale_factor),
+                             target_transform=target_transform(cropsize))
 
 
-def get_test_set(upscale_factor):
-    test_dir = "dataset/Urban100"
-    crop_size = calculate_valid_crop_size(256, upscale_factor)
+def get_test_set(dataset, crop_size, upscale_factor):
+    test_dir = join("dataset", dataset)
+    cropsize = calculate_valid_crop_size(crop_size, upscale_factor)
 
     return DatasetFromFolder(test_dir,
                              input_transform=input_transform(
-                                 crop_size, upscale_factor),
-                             target_transform=target_transform(crop_size))
+                                 cropsize, upscale_factor),
+                             target_transform=target_transform(cropsize))
diff --git a/main.py b/main.py
@@ -16,15 +16,18 @@
 from model import VDSR
 
 
-parser = argparse.ArgumentParser(
-    description='PyTorch VDSR')
+parser = argparse.ArgumentParser(description='PyTorch VDSR')
+parser.add_argument('--dataset', type=str, default='BSDS300',
+                    required=True, help="dataset directory name")
+parser.add_argument('--crop_size', type=int, default=256,
+                    required=True, help="network input size")
 parser.add_argument('--upscale_factor', type=int, default=2,
                     required=True, help="super resolution upscale factor")
-parser.add_argument('--batch_size', type=int, default=64,
-                    help='training batch size')
+parser.add_argument('--batch_size', type=int, default=128,
+                    help="training batch size")
 parser.add_argument('--test_batch_size', type=int,
-                    default=10, help='testing batch size')
-parser.add_argument('--epochs', type=int, default=100,
+                    default=32, help="testing batch size")
+parser.add_argument('--epochs', type=int, default=10,
                     help='number of epochs to train for')
 parser.add_argument('--lr', type=float, default=0.001,
                     help='Learning Rate. Default=0.001')
@@ -35,7 +38,7 @@
 parser.add_argument("--weight-decay", "--wd", default=1e-4,
                     type=float, help="Weight decay, Default: 1e-4")
 parser.add_argument('--cuda', action='store_true', help='use cuda?')
-parser.add_argument('--threads', type=int, default=16,
+parser.add_argument('--threads', type=int, default=128,
                     help='number of threads for data loader to use')
 parser.add_argument('--gpuids', default=[0], nargs='+',
                     help='GPU ID for using')
@@ -57,12 +60,15 @@ def main():
 
     if opt.cuda and not torch.cuda.is_available():
         raise Exception("No GPU found, please run without --cuda")
+
     cudnn.benchmark = True
 
-    train_set = get_training_set(
+    train_set = get_training_set(opt.dataset, opt.crop_size,
         opt.upscale_factor, opt.add_noise, opt.noise_std)
-    validation_set = get_validation_set(opt.upscale_factor)
-    test_set = get_test_set(opt.upscale_factor)
+    validation_set = get_validation_set(
+        opt.dataset, opt.crop_size, opt.upscale_factor)
+    test_set = get_test_set(
+        opt.dataset, opt.crop_size, opt.upscale_factor)
     training_data_loader = DataLoader(
         dataset=train_set, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True)
     validating_data_loader = DataLoader(
@@ -90,16 +96,34 @@ def main():
         start_time = time.time()
         test(model, criterion, testing_data_loader)
         elapsed_time = time.time() - start_time
-        print("===> average {:.2f} image/sec for processing".format(
+        print("===> average {:.2f} image/sec for test".format(
             100.0/elapsed_time))
         return
 
+    train_time = 0.0
+    validate_time = 0.0
     for epoch in range(1, opt.epochs + 1):
+        start_time = time.time()
         train(model, criterion, epoch, optimizer, training_data_loader)
+        elapsed_time = time.time() - start_time
+        train_time += elapsed_time
+        print("===> {:.2f} seconds to train this epoch".format(
+            elapsed_time))
+        start_time = time.time()
         validate(model, criterion, validating_data_loader)
+        elapsed_time = time.time() - start_time
+        validate_time += elapsed_time
+        print("===> {:.2f} seconds to validate this epoch".format(
+            elapsed_time))
         if epoch % 10 == 0:
             checkpoint(model, epoch)
 
+    print("===> average training time per epoch: {:.2f} seconds".format(train_time/opt.epochs))
+    print("===> average validation time per epoch: {:.2f} seconds".format(validate_time/opt.epochs))
+    print("===> training time: {:.2f} seconds".format(train_time))
+    print("===> validation time: {:.2f} seconds".format(validate_time))
+    print("===> total training time: {:.2f} seconds".format(train_time+validate_time))
+
 
 def adjust_learning_rate(epoch):
     """Sets the learning rate to the initial LR decayed by 10 every 10 epochs"""
@@ -185,4 +209,7 @@ def checkpoint(model, epoch):
 
 
 if __name__ == '__main__':
+    start_time = time.time()
     main()
+    elapsed_time = time.time() - start_time
+    print("===> total time: {:.2f} seconds".format(elapsed_time))
