Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: enable encoder-only architecture for nnU-Net #6

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

strasserpatrick
Copy link

@strasserpatrick strasserpatrick commented Jul 31, 2024

Hello, thanks for the great work.

I am exploring self-supervised pretraininig for nnU-Net. For that, I do encoder-only pretraining and then transfering the learned encoder weights to the final full U-Net architecture on the finetuning.

With this little adjustment, the workflow for that is quite simple.

  1. I follow the guide of nnU-Net for pretraining and finetuning of nnU-Net
  2. I edit the plans.json file for the pretraining configuration and change the network architecture to the PlainConvEncoder. With the kwargs that I add with this PR, the additional decoder configuration used for the finetuning gets simply ignored, no more plans file editing.
...
"architecture": {
                "network_class_name": "dynamic_network_architectures.building_blocks.plain_conv_encoder.PlainConvEncoder",
                "arch_kwargs": {
                    "n_stages": 6,
                    "features_per_stage": [
                        32,
                        64,
                        128,
                        256,
                        320,
                        320
                    ],
                    "conv_op": "torch.nn.modules.conv.Conv3d",
                    "kernel_sizes": [
...

When finished, I plan to do a PR on nnU-Net for the self-supervised learning, if you are interested in that :)

With the nnU-Net and its plans files, everything is nice and configurable. The additional kwargs allow me to quickly only initialize the encoder in nnUNetTrainers allowing me the framework for self-supervised pretraining which is all about training the feature-extractor (encoder) with pseudo-supervised tasks.

Let me know what you think of this

@strasserpatrick strasserpatrick marked this pull request as ready for review July 31, 2024 17:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant