About training details. ¶
By: wang719 on March 1, 2024, 4:18 a.m.
Hi, can you share more details about training the baseline model? Can I train the model with the exactly same hypermeter with single 3090Ti GPU?
By: MJJdG on March 3, 2024, 9:21 p.m.
Hi Wang, training the baseline model took two steps to also incorporate the partially annotated data for 3D segmentation.
For the first step, we trained a 3D residual encoder nnUnet using the fully annotated data only. We then used this model to predict 3D masks for the partially annotated data (DeepLesion & CCC18). We excluded from these predicted cases those where the long or short-axis diameter error was larger than 10% of the lesion size.
For the second step, we used the predicted 3D masks to pre-train a new 3D resenc nnUnet for 1000 epochs. We then fine-tuned this model for 500 epochs with a starting learning rate of 2.5e-3 on the fully annotated data. We trained these models using the full VOI, and without any resampling.
The plans file with the exact hyperparameters is available on zenodo. We trained our model with a batch size of 3 on a single A100 GPU with 40GB of VRAM. To train it on a 3090Ti you could try reducing the batch size or decreasing the size of the VOI.