brandondube
commented
3 years ago Could you provide an example of what is broken when you just replace prysm.mathops.engine.numpy with tensorflow or pytorch?
For what it's worth, in my opinion there are not inherent advantages to autodiff from tf/pytorch/etc. These libraries, generally, lack critical features for image based wavefront sensing. Pytorch lacked complex number support (a hard requirement) for a long time, and tensorflow takes longer to set up its tape machine than some reverse calculations take to run in their entirety. Often array insertion and slicing ops are forbidden as "nondifferentiable" which is true, but not strictly so (with valid priors for these problems, they are differentiable) Some "batch updates" in the adjoint pass cannot utilize work sharing from the forward pass because those libs are blind to some of the "oracle knowledge" a skilled implementer has. For example, this tiny amount of code for a monochromatic, single plane image based wavefront sensing( "phase retrieval") problem runs more than 5x faster than the equivalent from any of jax, pytorch, or tf:
class NLOptPRModel:
def __init__(self, amp, wvl, basis, data):
self.amp = amp
self.wvl = wvl
self.basis = basis
self.D = data
def update(self, x):
phs = np.tensordot(basis, x, axes=(0,0))
W = (2 * np.pi / self.wvl) * phs
g = self.amp * np.exp(1j * W)
G = ft_fwd(g)
I = np.abs(G)**2
E = np.sum((I-self.D)**2)
self.W = W
self.g = g
self.G = G
self.I = I
self.E = E
return
def fwd(self, x):
self.update(x)
return self.E
def rev(self, x):
self.update(x)
Ibar = 2 * (self.I - self.D)
Gbar = 2 * Ibar * self.G
gbar = ft_rev(Gbar)
Wbar = 2 * np.pi / self.wvl * np.imag(gbar * np.conj(self.g))
abar = np.tensordot(self.basis, Wbar)
self.Ibar = Ibar
self.Gbar = Gbar
self.gbar = gbar
self.Wbar = Wbar
self.abar = abar
return abar(This code is (C) me, you should attribute as such if you copy it)
You can convert it from a low dimensional modal estimator to an extremely high dimensional elementwise or "pixelwise" estimator by replacing phs = ... with phs = x.reshape(self.amp.shape), and return Wbar instead of abar. Obviously comment out the projection in the reverse pass, as well. For a high dimensional problem (~262k variables for 512x512) you will really want to prune all of the dead variables -- those outside the pupil's support --. 512x512 = 262,144, but at Q=2 there are only 51,445 pixels in the pupil -- removing 80% of the parameter space is a huge benefit. It's a numpy one liner (two liner if you value clarity) to do that.
Making the BFGS updates on a GPU will be faster for a large problem, but you may be surprised how big the array size is to make that worthwhile. L-BFGS-B from scipy does 512x512 single plane image based wavefront sensing in about 1.2s on a commodity desktop computer, with 65 iterations in that time. 18ms/iter is not that slow, and with one fwd and one rev eval per iter, <9ms each does not particularly motivate the GPU. (It would be a different story were it, say, 1s). It takes ~54 seconds to do 1000 iterations of the 262k variable problem (52ms/iter). I would posit you need something like a four, maybe 16 million variables to motivate the GPU. And then you must have some extremely high frequency phase features to solve, which will require a special diversity that will probably occupy more of your time to identify than the computation.
Anyway - in short, manual autograd, if you care about performance, trounces any library doing it for you.
usryokousha
I am doing some phase retrieval work for systems with large fields of view and typically I end up wanting to recover the phase for multiple local regions. In terms of use for astronomy and photography the use of tensors as opposed to 2D arrays could be very useful.
Tensorflow now includes conventional optimizers such as CG and LBFGS (tensorflow-probability). There are some inherent advantages to using the auto differentiation features offered by the library as well.