Pretrained weights for UNET and First stage.
bhosalems opened this issue · 4 comments
Hi authors,
Sorry for posting many question, but I am facing some issues starting the training. Specifically, after few steps my loss is nan, I was wondering this could be because of the checkpoints I have used to start the training.
For example, I have used UNET checkpoint as mentioned in #5 but, I am not sure what things you retained from the checkpoints. I see mention of removing class embedder weights, but apart from that did you also removed first stage weights? You do load your own first stage weights the model later, but because of sequence of operations I believe UNET weights are loaded after first stage, essentially overwriting the weights from your own first stage. Also, i am not sure what exactly did you change in the lightening config for UNET weights to work?
Similarly, I extracted the weights for the first stage. But i am unsure if I did that correctly.
I would appreciate if you could also share the weights for UNET and your first stage.
Thank you!
So to clarify, are these UNET weights the ones as provided by the pretrained LDM? Or are they already fine-tuned on the TCGA dataset?
Nevermind, from sampling it is quite clear that these have not been trained on pathological images yet