Fix default parameters

Author: alebeck

README.md

@@ -42,7 +42,7 @@ Run the training script via
 ```bash
 python run_train.py --dataset <...> --data-dir <...> --save-dir <...>
 ```
-You can see all available training options by running 
+Hyperparameter defaults are set to the values from the paper. Depending on the dataset, you have to adjust the number of epochs (`--num-epochs`) and the scheduler step size (`--lr-step`), see appendix A of the paper. You can see all available training options by running 
 ```bash
 python run_train.py -h
 ```

arguments/__init__.py

@@ -0,0 +1,2 @@
+from .train import parser as train_parser
+from .eval import parser as eval_parser

arguments/eval.py

@@ -0,0 +1,14 @@
+import configargparse
+
+parser = configargparse.ArgumentParser()
+parser.add_argument('-c', '--config', is_config_file=True, help='Path to the config file', type=str)
+
+parser.add_argument('--checkpoint', type=str, required=True, help='Checkpoint path to evaluate')
+parser.add_argument('--dataset', type=str, required=True, help='Name of the dataset')
+parser.add_argument('--data-dir', type=str, required=True, help='Root directory of the dataset')
+parser.add_argument('--num-workers', type=int, default=8, metavar='N', help='Number of dataloader worker processes')
+parser.add_argument('--batch-size', type=int, default=8)
+parser.add_argument('--crop-size', type=int, default=256, help='Size of the input (squared) patches')
+parser.add_argument('--scaling', type=int, default=8, help='Scaling factor')
+parser.add_argument('--in-memory', default=False, action='store_true', help='Hold data in memory during evaluation')
+parser.add_argument('--feature-extractor', type=str, default='UResNet', help='Feature extractor for edge potentials')

arguments/train.py

@@ -0,0 +1,43 @@
+import configargparse
+
+parser = configargparse.ArgumentParser()
+parser.add_argument('-c', '--config', is_config_file=True, help='Path to the config file', type=str)
+
+# general
+parser.add_argument('--save-dir', required=True, help='Path to directory where models and logs should be saved')
+parser.add_argument('--logstep-train', default=10, type=int, help='Training log interval in steps')
+parser.add_argument('--save-model', default='both', choices=['last', 'best', 'no', 'both'])
+parser.add_argument('--val-every-n-epochs', type=int, default=1, help='Validation interval in epochs')
+parser.add_argument('--resume', type=str, default=None, help='Checkpoint path to resume')
+parser.add_argument('--seed', type=int, default=12345, help='Random seed')
+parser.add_argument('--wandb', action='store_true', default=False, help='Use Weights & Biases instead of TensorBoard')
+parser.add_argument('--wandb-project', type=str, default='graph-sr', help='Wandb project name')
+
+# data
+parser.add_argument('--dataset', type=str, required=True, help='Name of the dataset')
+parser.add_argument('--data-dir', type=str, required=True, help='Root directory of the dataset')
+parser.add_argument('--num-workers', type=int, default=8, metavar='N', help='Number of dataloader worker processes')
+parser.add_argument('--batch-size', type=int, default=8)
+parser.add_argument('--crop-size', type=int, default=256, help='Size of the input (squared) patches')
+parser.add_argument('--scaling', type=int, default=8, help='Scaling factor')
+parser.add_argument('--max-rotation', type=float, default=15., help='Maximum rotation angle (degrees)')
+parser.add_argument('--no-flip', action='store_true', default=False, help='Switch off random flipping')
+parser.add_argument('--in-memory', action='store_true', default=False, help='Hold data in memory during training')
+
+# training
+parser.add_argument('--loss', default='l1', type=str, choices=['l1', 'mse'])
+parser.add_argument('--num-epochs', type=int, default=250)
+parser.add_argument('--optimizer', default='adam', choices=['sgd', 'adam'])
+parser.add_argument('--lr', type=float, default=0.0001)
+parser.add_argument('--momentum', type=float, default=0.9)
+parser.add_argument('--w-decay', type=float, default=1e-5)
+parser.add_argument('--lr-scheduler', type=str, default='step', choices=['no', 'step', 'plateau'])
+parser.add_argument('--lr-step', type=int, default=10, help='LR scheduler step size (epochs)')
+parser.add_argument('--lr-gamma', type=float, default=0.9, help='LR decay rate')
+parser.add_argument('--skip-first', action='store_true', help='Don\'t optimize during first epoch')
+parser.add_argument('--gradient-clip', type=float, default=0.01, help='If > 0, clips gradient norm to that value')
+
+# model
+parser.add_argument('--feature-extractor', type=str, default='UResNet', help='Feature extractor for edge potentials')
+parser.add_argument('--lambda-init', type=float, default=1., help='Graph lambda parameter initialization')
+parser.add_argument('--mu-init', type=float, default=.1, help='Graph mu parameter initialization')

model/graph_sr_net.py

@@ -46,7 +46,7 @@ def __init__(
             scaling: int,
             crop_size=256,
             feature_extractor='UResNet',
-            pretrained=False,
+            pretrained=True,
             lambda_init=1.0,
             mu_init=0.1
     ):

run_eval.py

@@ -3,28 +3,16 @@
 from collections import defaultdict
 import time
 
-import configargparse
 import torch
 from torchvision.transforms import Normalize
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from arguments import eval_parser
 from model import GraphSuperResolutionNet
 from data import MiddleburyDataset, NYUv2Dataset, DIMLDataset
 from utils import to_cuda
 
-parser = configargparse.ArgumentParser()
-parser.add_argument('-c', '--config', is_config_file=True, help='Path to the config file', type=str)
-parser.add_argument('--checkpoint', type=str, required=True, help='Checkpoint path to evaluate')
-parser.add_argument('--dataset', type=str, required=True, help='Name of the dataset')
-parser.add_argument('--data-dir', type=str, required=True, help='Root directory of the dataset')
-parser.add_argument('--num-workers', type=int, default=8, metavar='N', help='Number of dataloader worker processes')
-parser.add_argument('--batch-size', type=int, default=8)
-parser.add_argument('--crop-size', type=int, default=256, help='Size of the input (squared) patches')
-parser.add_argument('--scaling', type=int, default=8, help='Scaling factor')
-parser.add_argument('--in-memory', default=False, action='store_true', help='Hold data in memory during evaluation')
-parser.add_argument('--feature-extractor', type=str, default='UResNet', help='Feature extractor for edge potentials')
-
 
 class Evaluator:
 
@@ -96,8 +84,8 @@ def resume(self, path):
 
 
 if __name__ == '__main__':
-    args = parser.parse_args()
-    print(parser.format_values())
+    args = eval_parser.parse_args()
+    print(eval_parser.format_values())
 
     evaluator = Evaluator(args)
 

run_train.py

@@ -3,7 +3,6 @@
 from collections import defaultdict
 import time
 
-import configargparse
 import numpy as np
 import torch
 from torch import optim
@@ -13,54 +12,11 @@
 from torchvision.transforms import Normalize
 from tqdm import tqdm
 
+from arguments import train_parser
 from model import GraphSuperResolutionNet
 from data import MiddleburyDataset, NYUv2Dataset, DIMLDataset
 from utils import new_log, to_cuda, seed_all
 
-parser = configargparse.ArgumentParser()
-parser.add_argument('-c', '--config', is_config_file=True, help='Path to the config file', type=str)
-
-# general
-parser.add_argument('--save-dir', required=True, help='Path to directory where models and logs should be saved')
-parser.add_argument('--logstep-train', default=10, type=int, help='Training log interval in steps')
-parser.add_argument('--save-model', default='both', choices=['last', 'best', 'no', 'both'])
-parser.add_argument('--val-every-n-epochs', type=int, default=1, help='Validation interval in epochs')
-parser.add_argument('--resume', type=str, default=None, help='Checkpoint path to resume')
-parser.add_argument('--seed', type=int, default=12345, help='Random seed')
-parser.add_argument('--wandb', action='store_true', default=False, help='Use Weights & Biases instead of TensorBoard')
-parser.add_argument('--wandb-project', type=str, default='graph-sr', help='Wandb project name')
-
-# data
-parser.add_argument('--dataset', type=str, required=True, help='Name of the dataset')
-parser.add_argument('--data-dir', type=str, required=True, help='Root directory of the dataset')
-parser.add_argument('--num-workers', type=int, default=8, metavar='N', help='Number of dataloader worker processes')
-parser.add_argument('--batch-size', type=int, default=8)
-parser.add_argument('--crop-size', type=int, default=256, help='Size of the input (squared) patches')
-parser.add_argument('--scaling', type=int, default=8, help='Scaling factor')
-parser.add_argument('--max-rotation', type=float, default=15., help='Maximum rotation angle (degrees)')
-parser.add_argument('--no-flip', action='store_true', default=False, help='Switch off random flipping')
-parser.add_argument('--in-memory', action='store_true', default=False, help='Hold data in memory during training')
-
-# training
-parser.add_argument('--loss', default='l1', type=str, choices=['l1', 'mse'])
-parser.add_argument('--num-epochs', type=int, default=250)
-parser.add_argument('--optimizer', default='adam', choices=['sgd', 'adam'])
-parser.add_argument('--lr', type=float, default=0.0001)
-parser.add_argument('--momentum', type=float, default=0.9)
-parser.add_argument('--w-decay', type=float, default=1e-5)
-parser.add_argument('--lr-scheduler', type=str, default='step', choices=['no', 'step', 'plateau'])
-parser.add_argument('--lr-step', type=int, default=10, help='LR scheduler step size (epochs)')
-parser.add_argument('--lr-gamma', type=float, default=0.9, help='LR decay rate')
-parser.add_argument('--skip-first', action='store_true', help='Don\'t optimize during first epoch')
-parser.add_argument('--gradient-clip', type=float, default=0., help='If > 0, clips gradient norm to that value')
-
-# model
-parser.add_argument('--feature-extractor', type=str, default='UResNet', help='Feature extractor for edge potentials')
-parser.add_argument('--pretrained', action='store_true', help='Initialize feature extractor with weights '
-                                                              'pretrained on ImageNet')
-parser.add_argument('--lambda-init', type=float, default=1., help='Graph lambda parameter initialization')
-parser.add_argument('--mu-init', type=float, default=.1, help='Graph mu parameter initialization')
-
 
 class Trainer:
 
@@ -76,9 +32,8 @@ def __init__(self, args: argparse.Namespace):
             args.scaling,
             args.crop_size,
             args.feature_extractor,
-            args.pretrained,
-            args.lambda_init,
-            args.mu_init
+            lambda_init=args.lambda_init,
+            mu_init=args.mu_init
         )
         self.model.cuda()
 
@@ -272,8 +227,8 @@ def resume(self, path):
 
 
 if __name__ == '__main__':
-    args = parser.parse_args()
-    print(parser.format_values())
+    args = train_parser.parse_args()
+    print(train_parser.format_values())
 
     if args.wandb:
         import wandb