cellarium-ai / cellarium-ml

Distributed single-cell data analysis.
BSD 3-Clause "New" or "Revised" License
22 stars 3 forks source link

Auto-generate `example_input_array` property #224

Open ordabayevy opened 3 months ago

ordabayevy commented 3 months ago

LightningModule has example_input_array property which if set provides following benefits:

  1. Debugging: Print input output layer dimensions
  2. Logging: Add the computational graph to tensorboard
  3. UX: It can help to infer the keys of batch_keys dictionary and help the user to determine what arguments need to be provided to the transforms and the model.

The proposal is to have some automated way of generating example_input_array on supported transforms/models. This might be possible because we use specific argument names in our transforms and models (e.g. x_ng, total_mrna_umis_n, etc).