SGD optimizer does not work for the fancy_ptycho.py example

[yoshikisd]

Bug description

I am trying to swap out the Adam_optimize with SGD_optimize in the fancy_ptycho.py example script to ensure that my refactoring of SGD_optimize, at a minimum, preserves its original behavior. Regarding why I'm doing a refactoring, see PR #17 (comment).

The script will run one epoch with a 'nan' loss and subsequently crashes. Failure log shown at the bottom.

Script to reproduce

On the master branch of cdtools, run the following script

import cdtools
from matplotlib import pyplot as plt

filename = 'example_data/lab_ptycho_data.cxi'
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)

model = cdtools.models.FancyPtycho.from_dataset(
    dataset,
    n_modes=3, 
    oversampling=2, 
    probe_support_radius=120, 
    propagation_distance=5e-3, 
    units='mm', 
    obj_view_crop=-50, 
)

device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)

for loss in model.SGD_optimize(50, dataset):
    print(model.report())

OS and package info

• Python 3.12.2
• Pytorch 2.5.1
• CUDA 12.4

Failure log

Epoch   1 completed in 3.27 s with loss nan
Traceback (most recent call last):
  File "/cdtools_yoshikisd/cdtools/examples/fancy_ptycho.py", line 25, in <module>
    for loss in model.SGD_optimize(50, dataset):#, lr=0.02, batch_size=10):
  File "/cdtools_yoshikisd/cdtools/src/cdtools/models/base.py", line 492, in AD_optimize
    raise res
  File "/cdtools_yoshikisd/cdtools/src/cdtools/models/base.py", line 457, in target
    result_queue.put(run_epoch(stop_event))
                     ^^^^^^^^^^^^^^^^^^^^^
  File "/cdtools_yoshikisd/cdtools/src/cdtools/models/base.py", line 417, in run_epoch
    loss += optimizer.step(closure).detach().cpu().numpy()
            ^^^^^^^^^^^^^^^^^^^^^^^
  File "/.conda/envs/cdtools_testing/lib/python3.12/site-packages/torch/optim/optimizer.py", line 487, in wrapper
    out = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/.conda/envs/cdtools_testing/lib/python3.12/site-packages/torch/optim/optimizer.py", line 91, in _use_grad
    ret = func(self, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.conda/envs/cdtools_testing/lib/python3.12/site-packages/torch/optim/sgd.py", line 112, in step
    loss = closure()
           ^^^^^^^^^
  File "/cdtools_yoshikisd/cdtools/src/cdtools/models/base.py", line 395, in closure
    sim_patterns = self.forward(*inp)
                   ^^^^^^^^^^^^^^^^^^
  File "/cdtools_yoshikisd/cdtools/src/cdtools/models/base.py", line 94, in forward
    return self.measurement(self.forward_propagator(self.interaction(*args)))
                                                    ^^^^^^^^^^^^^^^^^^^^^^^
  File "/cdtools_yoshikisd/cdtools/src/cdtools/models/fancy_ptycho.py", line 502, in interaction
    exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/cdtools_yoshikisd/cdtools/src/cdtools/tools/interactions/interactions.py", line 472, in ptycho_2D_sinc
    output = shifted_probe * selections[...,None,:,:]
             ~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (278) must match the size of tensor b (0) at non-singleton dimension 3

[allevitan]

Looks to me like what's happening is the gradients in general are very high (leading eventually to NaNs), and specifically that the translation offsets are getting updated enough that they're eventually indexing regions outside of the object array (which causes the crash).

So I think the solution is to tune the learning rates (which, for SGD, probably need to be per-parameter) if we need an example script that uses SGD. It also could be a good idea to write an explicit catch for this error with a more readable error message (when the translation offsets lead to indexing a region outside of the object array) because it does happen reasonably often when something goes wrong and at present the error is not so easy to recognize.

[yoshikisd]

Ok that makes sense. I've confirmed that I can get a reconstruction to run if I...

• turn off grad_required for the translation offsets
• Make the learning rate extremely tiny (compared to those used in Adam_optimize)

Below is an example script with the changes:

    import cdtools
    from matplotlib import pyplot as plt

    filename = 'example_data/lab_ptycho_data.cxi' 
    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)

    model = cdtools.models.FancyPtycho.from_dataset(...)

+ model.translation_offsets.requires_grad = False

    device = 'cuda'
    model.to(device=device)
    dataset.get_as(device=device)

-  for loss in model.SGD_optimize(50, dataset):
+  for loss in model.SGD_optimize(50, dataset, lr=0.0000002):
        print(model.report())

Ideally, an example script would show how to tune these learning rates per-parameter. But for what I want to accomplish, this is good enough.

I also agree that we need to include a error message to catch issues that are (very likely) caused by high gradients.

For now, I'll close this issue since what I was seeing wasn't really a bug. I'll open up a branch for including the error catch. Perhaps it would be good to open an issue for creating an example for SGD. But I'm not really sure how much time we really want to invest in making the example since I'm not sure if there's a situation where we want to use SGD over Adam.