From a356f038a8d0a8d4cbe9dc1f54dcbd8989f50d57 Mon Sep 17 00:00:00 2001
From: Emilio Martinez <emilio.martinez@fing.edu.uy>
Date: Sun, 26 Mar 2023 12:21:06 -0300
Subject: [PATCH] Added epoch parameter for training and checkpoints

---
 KAIR/main_train_drunet.py      |  2 +-
 KAIR/options/train_drunet.json | 18 ++++++++++--------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/KAIR/main_train_drunet.py b/KAIR/main_train_drunet.py
index 33b720e..95a7fcb 100644
--- a/KAIR/main_train_drunet.py
+++ b/KAIR/main_train_drunet.py
@@ -159,7 +159,7 @@ def main(json_path='options/train_drunet.json'):
     # ----------------------------------------
     '''
 
-    for epoch in range(1000000):  # keep running
+    for epoch in range(opt['train']['epochs']):  # keep running
         if opt['dist']:
             train_sampler.set_epoch(epoch)
 
diff --git a/KAIR/options/train_drunet.json b/KAIR/options/train_drunet.json
index 4e1cc93..809f948 100644
--- a/KAIR/options/train_drunet.json
+++ b/KAIR/options/train_drunet.json
@@ -1,7 +1,7 @@
 {
   "task": "drunet"  //  root/task/images-models-options
   , "model": "plain" // "plain"
-  , "gpu_ids": []
+  , "gpu_ids": [0]
 
   , "scale": 1       // broadcast to "netG" if SISR
   , "n_channels": 2  // broadcast to "datasets", 1 for grayscale, 3 for color
@@ -18,7 +18,7 @@
     "train": {
       "name": "train_dataset"           // just name
       , "dataset_type": "ffdnet"         // "dncnn" | "dnpatch" for dncnn,  | "fdncnn" | "ffdnet" | "sr" | "srmd" | "dpsr" | "plain" | "plainpatch"
-      , "dataroot_H": "trainsets/trainH"// path of H training dataset
+      , "dataroot_H": "trainsets/web_images_train"// path of H training dataset
       , "dataroot_L": null              // path of L training dataset
       , "H_size": 128                    // patch size 40 | 64 | 96 | 128 | 192
       , "dataloader_shuffle": true
@@ -28,7 +28,7 @@
     , "test": {
       "name": "test_dataset"            // just name
       , "dataset_type": "ffdnet"         // "dncnn" | "dnpatch" for dncnn,  | "fdncnn" | "ffdnet" | "sr" | "srmd" | "dpsr" | "plain" | "plainpatch"
-      , "dataroot_H": "testsets/set12"  // path of H testing dataset
+      , "dataroot_H": "testsets/web_images_test"  // path of H testing dataset
       , "dataroot_L": null              // path of L testing dataset
     }
   }
@@ -52,7 +52,8 @@
   }
 
   , "train": {
-    "G_lossfn_type": "l1"               // "l1" preferred | "l2sum" | "l2" | "ssim" 
+      "epochs": 15                      // number of epochs to train
+    , "G_lossfn_type": "l1"               // "l1" preferred | "l2sum" | "l2" | "ssim" 
     , "G_lossfn_weight": 1.0            // default
     , "G_tvloss_weight": 1.0            // total variation weight
 
@@ -66,9 +67,10 @@
 
     , "G_regularizer_orthstep": null    // unused
     , "G_regularizer_clipstep": null    // unused
-
-    , "checkpoint_test": 5000           // for testing
-    , "checkpoint_save": 5000           // for saving model
-    , "checkpoint_print": 200           // for print
+    
+    // iteration (batch step) checkpoints
+    , "checkpoint_test": 60           // for testing
+    , "checkpoint_save": 60           // for saving model
+    , "checkpoint_print": 20           // for print
   }
 }