Merge pull request #1175 from computational-cell-analytics/dev

Dev

Author: constantinpape

micro_sam/__version__.py

@@ -1 +1 @@
-__version__ = "1.7.4"
+__version__ = "1.7.5"

micro_sam/instance_segmentation.py

@@ -789,16 +789,21 @@ def get_unetr(
         unetr_state_dict = unetr.state_dict()
         for k, v in unetr_state_dict.items():
             if not k.startswith("encoder"):
-                if flexible_load_checkpoint:  # Whether allow reinitalization of params, if not found.
+                # Whether allow reinitalization of params, if not found or mismatched.
+                if flexible_load_checkpoint:
                     if k in decoder_state:  # First check whether the key is available in the provided decoder state.
-                        unetr_state_dict[k] = decoder_state[k]
+                        if v.shape != decoder_state[k].shape:   # Then check if the sizes mismatch.
+                            warnings.warn(f"Shape of '{k}' did not match. Hence, we reinitialize it.")
+                            unetr_state_dict[k] = v
+                        else:
+                            unetr_state_dict[k] = decoder_state[k]
                     else:  # Otherwise, allow it to initialize it.
                         warnings.warn(f"Could not find '{k}' in the pretrained state dict. Hence, we reinitialize it.")
                         unetr_state_dict[k] = v
 
-                else:  # Whether be strict on finding the parameter in the decoder state.
+                else:  # Be strict on finding the parameter in the decoder state.
                     if k not in decoder_state:
-                        raise RuntimeError(f"The parameters for '{k}' could not be found.")
+                        raise RuntimeError(f"The parameters for '{k}' could not be found or has a size mismatch.")
                     unetr_state_dict[k] = decoder_state[k]
 
         unetr.load_state_dict(unetr_state_dict)
@@ -1143,6 +1148,12 @@ def generate(
         else:
             if halo is None:
                 raise ValueError("You must pass a value for halo if tile_shape is given.")
+
+            # Shards are not thread-safe for parallel writing! So if we have shards we have to use them for tiling.
+            # This is ok in terms efficiency as GPU tiles are small; shards should still be manegable for the watershed.
+            if isinstance(segmentation, zarr.Array) and getattr(segmentation, "shards", None) is not None:
+                tile_shape = segmentation.shards
+
             segmentation = _watershed_from_center_and_boundary_distances_parallel(
                 center_distances=self._center_distances,
                 boundary_distances=self._boundary_distances,

micro_sam/training/training.py

@@ -220,6 +220,7 @@ def train_sam(
     verify_n_labels_in_loader: Optional[int] = 50,
     box_distortion_factor: Optional[float] = 0.025,
     overwrite_training: bool = True,
+    strict_decoder_loading: bool = True,
     **model_kwargs,
 ) -> None:
     """Run training for a SAM model.
@@ -265,6 +266,10 @@ def train_sam(
         overwrite_training: Whether to overwrite the trained model stored at the same location.
             By default, overwrites the trained model at each run.
             If set to 'False', it will avoid retraining the model if the previous run was completed.
+        strict_decoder_loading: Whether to require that the pre-trained decoder in the checkpoint, if present,
+            exactly matches the instance segmentation decoder. Decoders may have a mismatch in the output
+            channels if they were pre-trained for a different task. If set to False, decoders with a different
+            output dimension can be loaded; the output channels will be re-initialized.
         model_kwargs: Additional keyword arguments for the `micro_sam.util.get_sam_model`.
     """
     with _filter_warnings(ignore_warnings):
@@ -294,7 +299,8 @@ def train_sam(
 
             # Get the UNETR.
             unetr = get_unetr(
-                image_encoder=model.sam.image_encoder, decoder_state=state.get("decoder_state", None), device=device,
+                image_encoder=model.sam.image_encoder, decoder_state=state.get("decoder_state", None),
+                device=device, flexible_load_checkpoint=not strict_decoder_loading,
             )
 
             # Get the parameters for SAM and the decoder from UNETR.
@@ -435,6 +441,7 @@ def train_instance_segmentation(
     peft_kwargs: Optional[Dict] = None,
     ignore_warnings: bool = True,
     overwrite_training: bool = True,
+    strict_decoder_loading: bool = True,
     **model_kwargs,
 ) -> None:
     """Train a UNETR for instance segmentation using the SAM encoder as backbone.
@@ -481,6 +488,10 @@ def train_instance_segmentation(
         overwrite_training: Whether to overwrite the trained model stored at the same location.
             By default, overwrites the trained model at each run.
             If set to 'False', it will avoid retraining the model if the previous run was completed.
+        strict_decoder_loading: Whether to require that the pre-trained decoder in the checkpoint, if present,
+            exactly matches the instance segmentation decoder. Decoders may have a mismatch in the output
+            channels if they were pre-trained for a different task. If set to False, decoders with a different
+            output dimension can be loaded; the output channels will be re-initialized.
         model_kwargs: Additional keyword arguments for the `micro_sam.util.get_sam_model`.
     """
 
@@ -498,7 +509,8 @@ def train_instance_segmentation(
         )
         device = get_device(device)
         model = get_unetr(
-            image_encoder=sam_model.sam.image_encoder, decoder_state=state.get("decoder_state", None), device=device,
+            image_encoder=sam_model.sam.image_encoder, decoder_state=state.get("decoder_state", None),
+            device=device, flexible_load_checkpoint=not strict_decoder_loading,
         )
 
         optimizer, scheduler = _get_optimizer_and_scheduler(