Skip to content

Add nnunet one-line trainer (DAT-897) - #67

Merged
luandalmazo merged 5 commits into
mainfrom
feat/DAT-897
Jun 17, 2026
Merged

Add nnunet one-line trainer (DAT-897)#67
luandalmazo merged 5 commits into
mainfrom
feat/DAT-897

Conversation

@luandalmazo

Copy link
Copy Markdown
Collaborator

Adds NNUNetTrainer, which runs the full nnUNet v2 pipeline (export -> fingerprint -> plan -> preprocess -> train -> predict -> import) in a single trainer.fit() call, logging per-epoch metrics to MLflow. Also adds the tutorial notebook.

Important to note:
The deploy bundle patches trainer_name = 'nnUNetTrainer' inside checkpoint_final.pth before logging to MLflow. This is basically a stopgap for DAT-903 (the deploy container cannot install datamint from PyPI when running a dev version). The patch is safe as long as _DatamintNNUNetTrainer does not override the network architecture. When we solve DAT-903 we should remove this patch and set nnUNet_extTrainer in load context, so users can modify the network if they want.

Closes DAT-897.

@Lucashsmello Lucashsmello left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll trust you on this PR. :)

Check test results with python3.10: https://github.com/SonanceAI/datamint-python-api/actions/runs/27279659681/job/80570314031?pr=67
I guess you used a feature of f-string not available on Python 3.10

@luandalmazo

Copy link
Copy Markdown
Collaborator Author

All reviewed and tested. I performed both unit and integration tests multiple times. Currently, the Git CI tests are returning errors due to the specific Python version being used and its particularities. Because of this, I have opened a new task (DAT-936) to add fallbacks and fix the tests.

@luandalmazo
luandalmazo merged commit 26d8588 into main Jun 17, 2026
2 of 3 checks passed
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.

2 participants