Autodiff wrt to trajectory #116
Conversation
|
@matthieutrs if you want to have a look ;) |
|
@paquiteau I handled the minor details.. locally my tests are passing. Ill add you to officially review it in free time. @alineyyy will handle your review comments. I shall try to see if I can support this in gpuNUFFT in meantime. Great job everyone! |
|
Ahhh I forgot, @alineyyy you need to do the same updates to finufft interface... |
There was a problem hiding this comment.
LGTM, Adding @paquiteau for some final checks and then we can merge when green.
src/mrinufft/operators/autodiff.py
Outdated
| (x,) = ctx.saved_tensors | ||
| return ctx.nufft_op.adj_op(dy), None | ||
| (x, traj) = ctx.saved_tensors | ||
|
|
There was a problem hiding this comment.
| grad_data = None | |
| grad_traj = None |
src/mrinufft/operators/autodiff.py
Outdated
| else: | ||
| grad_traj = None |
There was a problem hiding this comment.
| else: | |
| grad_traj = None |
| from ..base import FourierOperatorCPU | ||
| from mrinufft._utils import proper_trajectory | ||
| from mrinufft._utils import proper_trajectory, get_array_module | ||
|
|
||
|
|
||
| def get_fourier_matrix(ktraj, shape, dtype=np.complex64, normalize=False): | ||
| """Get the NDFT Fourier Matrix.""" |
There was a problem hiding this comment.
some documentation as well if possible ?
src/mrinufft/operators/base.py
Outdated
| if not AUTOGRAD_AVAILABLE or not self.autograd_available: | ||
| raise ValueError("Autograd not available, ensure torch is installed.") |
There was a problem hiding this comment.
| if not AUTOGRAD_AVAILABLE or not self.autograd_available: | |
| raise ValueError("Autograd not available, ensure torch is installed.") | |
| if not AUTOGRAD_AVAILABLE: | |
| raise ValueError("Autograd not available, ensure torch is installed.") | |
| if not self.autograd_available: | |
| raise ValueError("Backend does not support auto-differentiation.") |
| self.shape, | ||
| self.n_trans, | ||
| self.eps, | ||
| dtype="complex64" if self.samples.dtype == "float32" else "complex128", |
There was a problem hiding this comment.
you can use DTYPE_R2C dict for that
| matrix = ( | ||
| xp.exp(-2j * xp.pi * traj_grid).to(dtype).to(device).clone() | ||
| if xp.__name__ == "torch" | ||
| else (xp.exp(-2j * xp.pi * traj_grid, dtype=dtype)) | ||
| ) |
There was a problem hiding this comment.
| matrix = ( | |
| xp.exp(-2j * xp.pi * traj_grid).to(dtype).to(device).clone() | |
| if xp.__name__ == "torch" | |
| else (xp.exp(-2j * xp.pi * traj_grid, dtype=dtype)) | |
| ) | |
| matrix = xp.exp(-2j * xp.pi * traj_grid) | |
| if xp.__name__ == "torch": | |
| matrix.to(dtype=dtype, device=device, copy=True) | |
| ) |
There was a problem hiding this comment.
I think I have to use matrix = matrix.to(...) here, otherwise it will occur an error in the autograd function of pytorch.
|
If the failing gpunufft test is too strict causing problems, just ensure the gradients are close and maybe just make the test less strict |
This resolves #103