Update CellariumAnnDataDataModule and CellariumModule
#125
Conversation
ordabayevy
changed the title
CellariumAnnDataDataModule and CellariumModule
ordabayevy
force-pushed
the
module
branch
2 times, most recently
from
605f261 to
32cca79
Compare
ordabayevy
force-pushed
the
module
branch
8 times, most recently
from
95a2bca to
11c767f
Compare
cellarium/ml/core/module.py
Outdated
|
|
||
| _transforms, _model = transforms, model | ||
|
|
||
| transforms = [uninitialize_object(transform) for transform in _transforms] if _transforms is not None else None |
There was a problem hiding this comment.
Perhaps also write two lines of comment about why you're uninitializing and reinitializing?
There was a problem hiding this comment.
Added a comment above explaining the logic behind un-initializing and initializing.
There was a problem hiding this comment.
model and transforms are not attributes of self (yet)-- does uninitializing them have any effect on the behavior of self.save_hyperparameters?
| # In order to achieve this, we temporarily re-assign `dadc` to its un-initialized state | ||
| # and then call `save_hyperparameters` which will save these values as hparams. | ||
| # Then, we re-assign `dadc` back to its initialized state. | ||
| # `initialize_object` handles the case when the object was passed as a dictionary of class path and init args. | ||
| _dadc = dadc | ||
| dadc = uninitialize_object(_dadc) | ||
| self.save_hyperparameters(logger=False) |
There was a problem hiding this comment.
can this be done more neatly with a context manager construct? also, dadc not not an attrib of self yet-- how would uninitializing it on line 93 have an effect on self.save_hyperparameters call on line 94?
There was a problem hiding this comment.
Having a context manager is good idea, let me think about it more.
save_hyperparameters looks up the variable from the locals() in the current frame using an argument name from the function signature.
There was a problem hiding this comment.
Another approach is to have model and transforms args of CellariumModule as dicts of class name and init args and delegate their initialization to CellariumModule instead of LightningCLI. In fact, there is a configure_model method of LightningModule that is used to initialize large models with FSDP and DeepSpeed strategies (https://lightning.ai/docs/pytorch/stable/advanced/model_init.html#model-parallel-training-fsdp-and-deepspeed). We can just always use it to initialize all our models. WDYT?
ordabayevy
force-pushed
the
module
branch
4 times, most recently
from
6f0fbf9 to
7cba8f2
Compare
save_hyperparameterssave_hyperparameterspulls init args from thelocals(). Therefore, we can uninitializetorch.nn.Modules before callingsave_hyperparameters.torch.nn.Modules in the__init__method if they were provided asclass_pathandinit_args.configarg from theCellariumModulethat was used to store hyper-parametersCellariumAnnDataDataModulewhere bothDistributedAnnDataCollectionandAnnDatacan be passedIncrementalPCAorOnePassMeanVarStdoptimizer is not needed.