add DOG to ADAM

README.md

@@ -35,7 +35,7 @@ All in all, this results in: 50 views (TOF) -> 25 subsets, 128 views -> 16 subse
 
 ### 1) Educated Warm Start (main_EWS_SAGA.py)
 
-To reduce the time required to reach the minimiser, we want to start closer to the minimiser. A better initialisation could reduce the number of steps an iterative algorithm needs and thus reduce the time. To this end, we employ a neural network, to learn a suitable initial image. The network is a (small) 3D convolutional neural network. It takes as input the OSEM image, the (preconditioned) gradient of the likelihood and the (preconditioned) gradient of the prior. As the preconditioner we use the usuall choice taken in BSREM.  
+To reduce the time required to reach the minimiser, we want to start closer to the minimiser. A better initialisation could reduce the number of steps an iterative algorithm needs and thus reduce the time. To this end, we employ a neural network, to learn a suitable initial image. The network is a (small) 3D convolutional neural network. It takes as input the OSEM image, the (preconditioned) gradient of the likelihood and the (preconditioned) gradient of the prior. As the preconditioner we use the usuall choice taken in BSREM.  The network weights are available [HERE](https://drive.google.com/file/d/1RcuP74EVpmqXB2UGXJRuEKcv-Fpv6jWd/view?usp=sharing), please download and put in a folder **checkpoint/** in this repo. (TODO: write automatic download script)
 
 This has the obvious disadavantage that the could also do one epoch of SGD for the same cost. Maybe one epoch is SGD would actually result in a better initial value? I dont know. One could try a version only relying on the initial OSEM image. 
 

adam.py

@@ -32,6 +32,7 @@ def __init__(self, data, initial, relaxation_eta,
         Set the filter to `None` if you don't want any.
         '''
         super().__init__(**kwargs)
+        self.initial = initial.copy()
         self.x = initial.copy()
         self.data = data
         self.num_subsets = len(data)
@@ -59,6 +60,10 @@ def __init__(self, data, initial, relaxation_eta,
         self.g_sq = initial.get_uniform_copy(0)
         self.x_update = initial.get_uniform_copy(0)
 
+        # DOG parameters
+        self.max_distance = 0 
+        self.sum_gradient = 0  
+
     def subset_sensitivity(self, subset_num):
         raise NotImplementedError
 
@@ -97,12 +102,22 @@ def update(self):
             self.update_filter.apply(self.x_update)
 
         if self.iteration == 0:
-            self.initial_step_size = 4/(self.x_update.norm() + 1e-3) #min(max(1/(self.x_update.norm() + 1e-3), 0.05), 3.0)
+            #self.initial_step_size = 4/(self.x_update.norm() + 1e-3) #min(max(1/(self.x_update.norm() + 1e-3), 0.05), 3.0)
+            step_size = 4/(self.x_update.norm() + 1e-3) #min(max(1/(self.x_update.norm() + 1e-3), 0.05), 3.0)
             #print("Choose step size as: ", self.initial_step_size)
             print("update norm: ", 4/(self.x_update.norm() + 1e-3))
 
-        step_size = self.step_size()
-        #print("alpha = ", step_size)
+        distance = (self.x - self.initial).norm()
+        if distance > self.max_distance:
+            self.max_distance = distance 
+
+        self.sum_gradient += self.x_update.norm()**2
+
+        if self.iteration > 0:
+            step_size = 1.25 * self.max_distance / np.sqrt(self.sum_gradient)
+
+        #step_size = self.step_size()
+        print("step size = ", step_size)
 
         self.x.sapyb(1.0, self.x_update, step_size, out=self.x)
         #self.x += step_size * self.x_update

main_ADAM_SAGA.py

@@ -9,8 +9,8 @@
 """
 from cil.optimisation.algorithms import Algorithm
 from cil.optimisation.utilities import callbacks
-#from sirf.contrib.partitioner import partitioner
-from utils.partioner_function_no_obj import data_partition
+from sirf.contrib.partitioner import partitioner
+#from utils.partioner_function_no_obj import data_partition
 
 #from adam_saga import Adam
 from adam import Adam
@@ -50,7 +50,8 @@ def __init__(self, data,
         data.prior.set_up(data.OSEM_image)
         for f in obj_funs: # add prior evenly to every objective function
             f.set_prior(data.prior)
-
+
+        relaxation_eta = 0.02 
         super().__init__(data=data_sub, 
                          obj_funs=obj_funs,
                          initial=data.OSEM_image, 

petric.py

@@ -14,6 +14,7 @@
 Options:
   --log LEVEL  : Set logging level (DEBUG, [default: INFO], WARNING, ERROR, CRITICAL)
 """
+
 import csv
 import logging
 import os
@@ -22,6 +23,9 @@
 from time import time
 from traceback import print_exc
 
+import torch 
+torch.cuda.set_per_process_memory_fraction(0.6)
+
 import numpy as np
 from skimage.metrics import mean_squared_error as mse
 from tensorboardX import SummaryWriter
@@ -33,13 +37,12 @@
 
 from datetime import datetime
 
-#import torch 
-#torch.cuda.set_per_process_memory_fraction(0.6)
+
 
 log = logging.getLogger('petric')
 TEAM = os.getenv("GITHUB_REPOSITORY", "SyneRBI/PETRIC-").split("/PETRIC-", 1)[-1]
 VERSION = os.getenv("GITHUB_REF_NAME", "")
-OUTDIR = Path(f"/o/logs/{TEAM}/{VERSION}" if TEAM and VERSION else "./output/" + datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
+OUTDIR = Path(f"/o/logs/{TEAM}/{VERSION}" if TEAM and VERSION else "./output/" + "ADAM_" + datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
 if not (SRCDIR := Path("/mnt/share/petric")).is_dir():
     SRCDIR = Path("./data")
 
@@ -247,14 +250,16 @@ def get_image(fname):
 if SRCDIR.is_dir():
     # create list of existing data
     # NB: `MetricsWithTimeout` initialises `SaveIters` which creates `outdir`
-    data_dirs_metrics = [(SRCDIR / "NeuroLF_Hoffman_Dataset", OUTDIR / "NeuroLF_Hoffman",
-                          [MetricsWithTimeout(outdir=OUTDIR / "NeuroLF_Hoffman", transverse_slice=72)]),
-                          (SRCDIR / "Siemens_mMR_NEMA_IQ", OUTDIR / "mMR_NEMA",
+    data_dirs_metrics = [   (SRCDIR / "Siemens_mMR_NEMA_IQ", OUTDIR / "mMR_NEMA",
                           [MetricsWithTimeout(outdir=OUTDIR / "mMR_NEMA", transverse_slice=72, coronal_slice=109)]),
-                         (SRCDIR / "Siemens_Vision600_thorax", OUTDIR / "Vision600_thorax",
-                          [MetricsWithTimeout(outdir=OUTDIR / "Vision600_thorax")]),
+                            (SRCDIR / "NeuroLF_Hoffman_Dataset", OUTDIR / "NeuroLF_Hoffman",
+                          [MetricsWithTimeout(outdir=OUTDIR / "NeuroLF_Hoffman", transverse_slice=72)]),
                           (SRCDIR / "Mediso_NEMA_IQ", OUTDIR / "Mediso_NEMA",
                             [MetricsWithTimeout(outdir=OUTDIR / "Mediso_NEMA")]),
+                         (SRCDIR / "Siemens_Vision600_thorax", OUTDIR / "Vision600_thorax",
+                          [MetricsWithTimeout(outdir=OUTDIR / "Vision600_thorax")]),
+                        (SRCDIR / "Siemens_mMR_NEMA_IQ_lowcounts", OUTDIR / "mMR_NEMA_lowcounts",
+                          [MetricsWithTimeout(outdir=OUTDIR / "mMR_NEMA_lowcounts")]),
                         (SRCDIR / "Siemens_mMR_ACR", OUTDIR / "Siemens_mMR_ACR",
                      [MetricsWithTimeout(outdir=OUTDIR / "Siemens_mMR_ACR")])]
 
@@ -276,7 +281,7 @@ def get_image(fname):
     from docopt import docopt
     args = docopt(__doc__)
     logging.basicConfig(level=getattr(logging, args["--log"].upper()))
-    from main_ADAM import Submission, submission_callbacks
+    from main_ADAM_SAGA import Submission, submission_callbacks
     assert issubclass(Submission, Algorithm)
     for srcdir, outdir, metrics in data_dirs_metrics:
         data = get_data(srcdir=srcdir, outdir=outdir)