Generate `out_dim` dynamically

[sfluegel05]

I was wondering if we could explore dynamically setting the out_dim parameter in the setup method of the LightningModule , leveraging the fact that it inherits DataHooks. This way, the model could fetch the number of labels (num_of_labels) directly from the data module after it's loaded.

By doing this, we could avoid requiring the out_dim parameter to be explicitly passed in the command line every time. Instead, the model could dynamically adapt based on the dataset being used. This would simplify the workflow and reduce potential user errors.

Originally posted by @aditya0by0 in #64 (comment)

[sfluegel05]

I think this is something worth investigating. It might that there are some constraints in the architecture that don't allow a dynamic calculation of out_dim, but I am not sure about that. @aditya0by0 If you have some time on your hands, please have a look and maybe try to implement a solution.

[aditya0by0]

• I have started a github discussion for an approach related to this issue, please refer: Dynamically Setting out_dim Based on LightningDataModule in LightningCLI Lightning-AI/pytorch-lightning#20602

[aditya0by0]

The prepare_data and setup_data methods have to be called at the end of initialization of most derived subclass, this is because certain parameters in subclass might change data preprocessing. For that custom python metaclass is needed to customize object creation/call. But custom metaclass interferes with lightning's save_hyperparameters method flow due to a bug in lightning.

For that a issue has been raised.

• Issue: KeyError: 'self' in save_hyperparameters() when custom metaclass used Lightning-AI/pytorch-lightning#20697

• Discussion: KeyError: 'self' in save_hyperparameters() when custom metaclass used Lightning-AI/pytorch-lightning#20693

• Stackoverflow: https://stackoverflow.com/questions/79554986/keyerror-self-in-save-hyperparameters-when-custom-metaclass-used-pytorch

• A similar issue which existed but different error origin: KeyError in save_hyperparameters while using in a subclass Lightning-AI/pytorch-lightning#18405

[aditya0by0]

@sfluegel05,
I’ve implemented a simpler CLI-side approach to solve this issue. It’s more of a minimal workaround that avoids the complexity of previous solutions and keeps the codebase concise.

The core idea is to call the data module’s methods (like prepare_data() and setup()) before linking its arguments to the model through compute_fn of any first link_arguments. This allows the necessary data preparation to be done upfront before any linking is performed.

If this behavior is not needed in the future, the added lines in CLI can be easily commented out to revert to the previous behavior—so it’s a non-intrusive and reversible change.

Please have a look at the PR and let me know your thoughts!

[aditya0by0]

The prepare_data and setup_data methods have to be called at the end of initialization of most derived subclass, this is because certain parameters in subclass might change data preprocessing. For that custom python metaclass is needed to customize object creation/call. But custom metaclass interferes with lightning's save_hyperparameters method flow due to a bug in lightning.

For that a issue has been raised.

• Issue: KeyError: 'self' in save_hyperparameters() when custom metaclass used Lightning-AI/pytorch-lightning#20697
• Discussion: KeyError: 'self' in save_hyperparameters() when custom metaclass used Lightning-AI/pytorch-lightning#20693
• Stackoverflow: https://stackoverflow.com/questions/79554986/keyerror-self-in-save-hyperparameters-when-custom-metaclass-used-pytorch
• A similar issue which existed but different error origin: KeyError in save_hyperparameters while using in a subclass Lightning-AI/pytorch-lightning#18405

For the above issue, the fix is raised in the below PR. Hence, after the fix is merged, we can use the metaclass approach instead of the argument linking in new lightning version. This could be one approach to resolve: #124

Lightning-AI/pytorch-lightning#21227