Streamlined training with new Trainer class (#77)

* move utility function outside of script * New trainer class for training reconstruction * Update docstring * Update changelog * Update to trainer save * Fix partial mask support bug * Fix docstrings. * Fix APGD rendering. --------- Co-authored-by: Eric Bezzam <[email protected]>
LCAV · Aug 29, 2023 · 58f747a · 58f747a
1 parent bba42c0
commit 58f747a
Show file tree

Hide file tree

Showing 10 changed files with 453 additions and 233 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -27,6 +27,7 @@ Added
 - New dataset for pair of original image and their measurement from a screen. See ``utils.dataset.MeasuredDataset`` and ``utils.dataset.MeasuredDatasetSimulatedOriginal``.
 - Support for unrolled loading and inference in the script ``admm.py``.
 - Tikhonov reconstruction for coded aperture measurements (MLS / MURA).
+- New ``Trainer`` class to train ``TrainableReconstructionAlgorithm`` with PyTorch.
 
 
 Changed

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -5,4 +5,5 @@ numpy>=1.22      # so that default dtype are correctly rendered
 torch>=1.10
 torchvision>=0.15.2
 torchmetrics>=0.11.4
+pyFFS>=2.2.3   # for waveprop
 waveprop>=0.0.5
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -21,7 +21,9 @@
     "torchmetrics.image",
     "scipy.ndimage",
     "pycsou.abc",
+    "pycsou.operator",
     "pycsou.operator.func",
+    "pycsou.operator.linop",
     "pycsou.opt.solver",
     "pycsou.opt.stop",
     "pycsou.runtime",
@@ -33,6 +35,8 @@
     "paramiko",
     "paramiko.ssh_exception",
     "perlin_numpy",
+    "hydra",
+    "hydra.utils",
     "scipy.special",
     "matplotlib.cm",
     "pyffs",

diff --git a/docs/source/reconstruction.rst b/docs/source/reconstruction.rst
@@ -55,7 +55,7 @@
    Accelerated Proximal Gradient Descent (APGD)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-   .. autoclass:: lensless.APGD
+   .. autoclass:: lensless.recon.apgd.APGD
       :special-members: __init__
 
 
@@ -88,4 +88,22 @@
    .. autoclass:: lensless.UnrolledADMM
       :members: batch_call
       :special-members: __init__
-      :show-inheritance:
+      :show-inheritance:
+
+
+   Reconstruction Utilities
+   ------------------------
+
+   .. autoclass:: lensless.recon.utils.Trainer
+      :members: 
+      :special-members: __init__
+
+   .. autofunction:: lensless.recon.utils.load_drunet
+
+   .. autofunction:: lensless.recon.utils.apply_denoiser
+
+   .. autofunction:: lensless.recon.utils.get_drunet_function
+
+   .. autofunction:: lensless.recon.utils.measure_gradient
+
+   .. autofunction:: lensless.recon.utils.create_process_network
diff --git a/lensless/eval/benchmark.py b/lensless/eval/benchmark.py
@@ -10,8 +10,6 @@
 from lensless.utils.dataset import DiffuserCamTestDataset
 from tqdm import tqdm
 
-from lensless.utils.io import load_image
-
 try:
     import torch
     from torch.utils.data import DataLoader