|
33 | 33 | from fl4health.utils.nnunet_utils import ( |
34 | 34 | NNUNET_DEFAULT_NP, |
35 | 35 | NNUNET_N_SPATIAL_DIMS, |
| 36 | + LocalPolyLRScheduler, |
36 | 37 | Module2LossWrapper, |
37 | 38 | NnunetConfig, |
38 | 39 | NnUNetDataLoaderWrapper, |
|
50 | 51 | # silences a bunch of deprecation warnings related to scipy.ndimage |
51 | 52 | # Raised an issue with nnunet. https://github.com/MIC-DKFZ/nnUNet/issues/2370 |
52 | 53 | warnings.filterwarnings("ignore", category=DeprecationWarning) |
| 54 | + import nnunetv2.training.nnUNetTrainer.nnUNetTrainer as nnUNetTrainerModule |
53 | 55 | from batchgenerators.utilities.file_and_folder_operations import load_json, save_json |
54 | 56 | from nnunetv2.experiment_planning.experiment_planners.default_experiment_planner import ExperimentPlanner |
55 | 57 | from nnunetv2.experiment_planning.plan_and_preprocess_api import extract_fingerprints, preprocess_dataset |
|
58 | 60 | from nnunetv2.training.dataloading.utils import unpack_dataset |
59 | 61 | from nnunetv2.training.loss.deep_supervision import DeepSupervisionWrapper |
60 | 62 | from nnunetv2.training.lr_scheduler.polylr import PolyLRScheduler |
61 | | - from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer |
62 | 63 | from nnunetv2.utilities.dataset_name_id_conversion import convert_id_to_dataset_name |
63 | 64 |
|
64 | 65 | # grpcio currently has a log spamming bug that seems to be triggered by multithreading/multiprocessing |
@@ -86,7 +87,7 @@ def __init__( |
86 | 87 | checkpoint_and_state_module: ClientCheckpointAndStateModule | None = None, |
87 | 88 | reporters: Sequence[BaseReporter] | None = None, |
88 | 89 | client_name: str | None = None, |
89 | | - nnunet_trainer_class: type[nnUNetTrainer] = nnUNetTrainer, |
| 90 | + nnunet_trainer_class: type[nnUNetTrainerModule.nnUNetTrainer] = nnUNetTrainerModule.nnUNetTrainer, |
90 | 91 | nnunet_trainer_class_kwargs: dict[str, Any] | None = None, |
91 | 92 | ) -> None: |
92 | 93 | """ |
@@ -191,7 +192,7 @@ def __init__( |
191 | 192 | self.nnunet_trainer_class_kwargs = ( |
192 | 193 | nnunet_trainer_class_kwargs if nnunet_trainer_class_kwargs is not None else {} |
193 | 194 | ) |
194 | | - self.nnunet_trainer: nnUNetTrainer |
| 195 | + self.nnunet_trainer: nnUNetTrainerModule.nnUNetTrainer |
195 | 196 | self.nnunet_config: NnunetConfig |
196 | 197 | self.plans: dict[str, Any] | None = None |
197 | 198 | self.steps_per_round: int # N steps per server round |
@@ -585,6 +586,12 @@ def setup_client(self, config: Config) -> None: |
585 | 586 | device=self.device, |
586 | 587 | **self.nnunet_trainer_class_kwargs, |
587 | 588 | ) |
| 589 | + |
| 590 | + # NOTE: Monkey Patch to force the nnunet_trainer to use our version of the PolyLRScheduler instead of |
| 591 | + # NnUnet version. This is because NnUnet's hasn't updated their scheduler to the new torch signature and |
| 592 | + # does not appear to intend to do so. The fix is very minimum. So we patch it here. |
| 593 | + nnUNetTrainerModule.PolyLRScheduler = LocalPolyLRScheduler |
| 594 | + |
588 | 595 | # nnunet_trainer initialization |
589 | 596 | self.nnunet_trainer.initialize() |
590 | 597 | # This is done by nnunet_trainer in self.on_train_start, we |
|
0 commit comments