Selecting Optimal Training Epochs for Cell2location

Hello everyone,

I would like to express my gratitude for the development of the Cell2location tool, which has proven to be invaluable for spatial deconvolution. I have a question related to training the model and would appreciate any insights from your experience.

In the tutorial for the “Cell2location: spatial mapping” section, the model is trained with “max_epochs=30000.” This seems like an extensive training process, which I understand might be necessary given that the batch size is set to the size of the entire dataset. However, I’m curious if 30000 is the recommended number of epochs as a heuristic, or if there is more flexibility.

Since I plan to experiment with the other two hyperparameters (N, alpha), training for 30000 epochs each time could be time-consuming. Moreover, I have noticed that the loss in the ELBO plot plateaus after around 20000 epochs.

My current plan is to train the model using 30000 epochs for the initial configuration and select an optimal epoch based on the ELBO plot. I would then use this chosen epoch for hyperparameter selection and data analysis throughout the remainder of the project. Do you think this approach is reasonable? If possible, could you share your experiences or recommendations for choosing the number of training epochs?

Thank you in advance for your assistance!