MrGiovanni / SuPreM

[ICLR 2024 Oral] Supervised Pre-Trained 3D Models for Medical Image Analysis (9,262 CT volumes + 25 annotated classes)
https://www.cs.jhu.edu/~alanlab/Pubs23/li2023suprem.pdf
Other
246 stars 8 forks source link

Inconsistency in `SegResNet` parameter count #8

Closed surajpaib closed 7 months ago

surajpaib commented 7 months ago

Hi @MrGiovanni,

Thank you for open-sourcing all the models and weights from this very insightful study.

I'm comparing some of our trained SegResNet architectures against the SuPreM SegResNet. I was surprised to see that the param count is 470M given that our SegResNet (with more down blocks and init_filters only amounts to about 87M params).

Upon inspecting the model definition shown here, https://github.com/MrGiovanni/SuPreM/blob/65e6cdd473857a0964b294234b1e870b5d538efa/target_applications/totalsegmentator/test.py#L135

I see that the parameter count is 4.7M and not 470M.

image

Could you clarify if I'm doing something wrong here or if there was a mistake in the mentioned param. count?

WenxuanChelsea commented 7 months ago

Hi @surajpaib

Thanks for pointing out the parameter count issue with SegResNet. The parameter count for SegResNet should be 4.7M. We have already updated it in GitHub and will ensure this update is included in the upcoming ICLR camera-ready version.

Best, Wenxuan

surajpaib commented 7 months ago

Thanks, @WenxuanChelsea, for your prompt clarification.

Do you have any insights on the selection of this particular configuration for the SegResNet?

I tried fine-tuning my task using your weights (keeping the encoder frozen and tuning only the decoder) and the U-Net performed far superior to the SegResNet. Do you think that a bigger SegResNet would reduce this difference?

WenxuanChelsea commented 7 months ago

Hi @surajpaib

We trained a SegResNet because this architecture won several MICCAI 2023 segmentation challenges, including KiTS23, BraTS23, etc, last year.

Your findings are consistent with what we observed in Table 2. Perhaps, the lower performance of SegResNet can be due to the tiny version (init_filters=16) that we chose by default. We suggest a larger SegResNet model, e.g., init_filters=32 for an improved performance.

Every year, hundreds of new backbones are developed, and it is hard to determine which one is the best (of which specific configuration works optimally). Therefore we are inviting research groups worldwide to help faithfully benchmark their first/last authored segmentation backbones using our AbdomenAtlas dataset. This has attracted over 50 teams so far, including U-Net, Transformer U-Net, Mamba U-Net, NAS, and their variants. We aim to present a more comprehensive benchmark in the near future.

If you are also interested in contributing to this large-scale benchmark, please contact Zongwei at zzhou82@jh.edu

Thanks, Wenxuan

surajpaib commented 7 months ago

Thank you! Closing this issue. We will reach out to Zongwei!