From bdf2b2636cde42e9674d2e85cb3b15b403de69ce Mon Sep 17 00:00:00 2001 From: alexdenker Date: Mon, 23 Sep 2024 14:52:53 +0100 Subject: [PATCH] add DOG to ADAM --- README.md | 2 +- adam.py | 21 ++++++++++++++++++--- main_ADAM_SAGA.py | 7 ++++--- petric.py | 23 ++++++++++++++--------- 4 files changed, 37 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 44bf0ed..b645730 100644 --- a/README.md +++ b/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. diff --git a/adam.py b/adam.py index 4b65937..d937d13 100644 --- a/adam.py +++ b/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 diff --git a/main_ADAM_SAGA.py b/main_ADAM_SAGA.py index 9ca139f..b303fc6 100644 --- a/main_ADAM_SAGA.py +++ b/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, diff --git a/petric.py b/petric.py index 8617492..62fa94e 100755 --- a/petric.py +++ b/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)