Repository navigation
SGD optimizer does not work for the fancy_ptycho.py example #33
Description
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
Activity
allevitan commentedon Jun 12, 2025
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 commentedon Jun 12, 2025
Ok that makes sense. I've confirmed that I can get a reconstruction to run if I...
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.