update parse_args

f1052063 · Aug 20, 2022 · 8cd031e · 8cd031e
1 parent 63f751b
commit 8cd031e
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 32 deletions.
diff --git a/.idea/.gitignore b/.idea/.gitignore
diff --git a/.idea/NTSCC_JSAC22.iml b/.idea/NTSCC_JSAC22.iml
diff --git a/.idea/deployment.xml b/.idea/deployment.xml
diff --git a/.idea/remote-mappings.xml b/.idea/remote-mappings.xml
diff --git a/data/datasets.py b/data/datasets.py
@@ -58,6 +58,14 @@ def worker_init_fn_seed(worker_id):
 
     return train_loader, test_loader
 
+def get_test_loader(config):
+    test_dataset = Datasets(config)
+    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
+                                              batch_size=1,
+                                              shuffle=False)
+
+    return test_loader
+
 
 def get_test_loader(config):
     test_dataset = Datasets(config)

diff --git a/main.py b/main.py
@@ -7,7 +7,7 @@
 from net.NTSCC_Hyperior import NTSCC_Hyperprior
 import torch.optim as optim
 from utils import *
-from data.datasets import get_loader
+from data.datasets import get_loader, get_test_loader
 from config import config
 
 
@@ -188,11 +188,13 @@ def main(argv):
     net = NTSCC_Hyperprior(config).cuda()
     model_path = args.checkpoint
     load_weights(net, model_path)
-    train_loader, test_loader = get_loader(config)
+
 
     if args.phase == 'test':
+        test_loader = get_test_loader(config)
         test(net, test_loader, logger)
     elif args.phase == 'train':
+        train_loader, test_loader = get_loader(config)
         global global_step
         G_params = set(p for n, p in net.named_parameters() if not n.endswith(".quantiles"))
         aux_params = set(p for n, p in net.named_parameters() if n.endswith(".quantiles"))