Per-parameter Learning Rate (PPLR) Implementation, Tensor Decomposition Models, and Tomography Updates

[cedriclim1]

What does this PR do?

First pass implementation of KPlanes implementation (https://brentyi.github.io/tilted/, https://sarafridov.github.io/K-Planes/). These tensor decomposition models require multiple parmeter optimization which is not handled in OptimizerMixin.

The solution @arthurmccray and I came up with is to setup a per-parameter learning rate PPLR abstract class which will be instantiated with models that have multiple parameters to be optimized to. This is implemented incore/ml/models/model_base.py. The most important part about this class is the get_params method, which is used for parsing the trainable parameters (see ObjectTensorDecomp's get_optimization_parameters) .

class PPLR(ABC):
    """
    Abstract base class for models that require multi-scale parameter optimization.
    """

    @abstractmethod
    def get_params(self) -> Dict[str, list[nn.Parameter]]:
        """
        Return a dictionary of parameters grouped by key.

        For example if your nn.Module has multiple optimizable parameter groups,
        you can return a dictionary with the keys "grids" and "sigma_net"
        (KPlanes example).
        """
        pass

The KPlanes and TILTED implementations are in core/ml/models in kplanes.py with KPlanes, KPlanesTILTED, and CPTilted. Note CPTIlted is used for the two-phase warmup that the paper recommends by pretraining the SO3 rotations of the planes in a lower representation space. so3params.py also has both quaternion and R9+SVD implementations which can be swapped as a parameter in KPlanesTILTED.

In tomography/object_models.py there is now a new object that handles the new tensor decomposition methods, ObjectTensorDecomp that inherits from ObjectINR and the following overloads are needed:

• optimizer_params property: The setter is changed, since we're building a a dictionary of parameters into our optimizers per https://docs.pytorch.org/docs/stable/optim.html#per-parameter-options.
• reconnect_optimizer_to_parameters: This one is the most annoying thing to implement, since now our parameters is a dictionary.

There are a few optional overrides that I've also included preemptively just incase if I want to add additional soft/hard constraints to these methods as you can see both apply_soft_constraints and apply_hard_constraints are still there.

Key change to OptimizerMixin: I have changed set_optimizer to handle a dictionary of parameters. I think this is better than my original idea of just implementing this in object_models.py since you would have to overload that function again. I think this is fine and still works for single optimization stuff.

There are also small changes to different aspects of the Tomography module, but mostly contained in object_models.py. Listing out the changes and rationale in the tomography files:

• tomography.py: I've disabled torch.bfloat16 autocasting since F.grid_sample does not have autograd implemented errors. I've also put the TILTED tensor decomp pretraining check by looking at the convergence of the SO3 rotations. I don't like this being here, but I'm not quite sure where else to put this.
• tomography_base.py: Type-hinting change for setting up distributed.
• tomography_opt.py: Parsing is already handled in object_models.py.

What the reviewer should do

Attached is a notebook that tests tensor decomposition methods in tomography: 0415_PPLR_testing.ipynb. The dataset here is from quantem-tutorials, so just set the directory of the same dataset to this notebook. The runtime can be increased by increasing the batch size, but the learning rates have to be scaled. Setting the batch size to 4096 seems to be fine? So feel free to change it and scale the learning rates by 2.

• Double check implementations across the board, especially the optimizer_mixin.py, and object_models.py
• Organization makes sense for tensor decomposition stuff. I'm actually not very happy with how core/ml is looking. Where do we draw the line of what lives in models (I think everything model related should live in there tbh: INR, CNN, etc...)?
• Check reconstruction workflows are still working i.e, ptychography due to the OptimizerMixin change.

[bobleesj]

@cedriclim1 the default VS code setting has been removed, so please pull from upstreadm/dev

[arthurmccray]

overall looks good! I made the necessary few changes to ptycho so it works with this, and overall i'm happy with it. Hopefully we're one step closer to not having to mess with OptimizerMixin ever again...

As to /core/ml organization, yes cnn.py, cnn_dense.py, inr.py, etc. should all be in /ml/models. I don't think it should be that bad to move them actually? An automatic refactor should handle everything in quantem fine, and they'll be accessible at the same level of the namespace (we still want to be able to from quantem.core.ml import CNN2d), so it shouldn't affeect the tutorials either really. If you move them (and whatever else you think) we can just re-run all the tutorials and tests to make sure nothing breaks.

[cedriclim1]

@arthurmccray ready for review again - addressed most comments, but some need follow-up.

[arthurmccray]

Overall looks good, just one or two things I saw. Did I miss anything that you needed my followup on?

[arthurmccray]

bump. slight preference to having static methods rotmat_to_r9 and r9_to_rotmat (the latter of which is simply called by as_matrix), but just having it in a classmethod would be fine too.

[arthurmccray]

please include here what these actually mean tho. Like volume is pretty self explanatory (presumably it's the full reconstructed volume), and coords would be for INR (is it the coords for the full volume? or just part of it), but pred and all_densities and obj should be explained in a line

[arthurmccray]

it probably would, but it should be easy for me to fix... but on the other hand it feels a little silly for it to be {"params": OptimizerType} for single Optimizers. Up to you!

[arthurmccray]

in this case i don't think it makes any difference as the optimizer_type property just directly returns _optimizer_type. But in theory it might not, and that's more why I think about using the properties internally. idk if that's officially good or bad practice though--in my comment i was honestly asking it as a question 😅