Disable WANDB for now + log param tweaks

Author: C-Achard

napari_cellseg3d/code_models/worker_training.py

@@ -68,8 +68,6 @@
 logger.debug(f"PRETRAINED WEIGHT DIR LOCATION : {PRETRAINED_WEIGHTS_DIR}")
 
 try:
-    import wandb
-
     WANDB_INSTALLED = True
 except ImportError:
     logger.warning(
@@ -411,25 +409,25 @@ def log_parameters(self):
         self.log(f"Using {self.config.num_classes} classes")
         self.log(f"Weight decay: {self.config.weight_decay}")
         self.log("* NCuts : ")
-        self.log(f"- Insensity sigma {self.config.intensity_sigma}")
+        self.log(f"- Intensity sigma {self.config.intensity_sigma}")
         self.log(f"- Spatial sigma {self.config.spatial_sigma}")
         self.log(f"- Radius : {self.config.radius}")
         self.log(f"* Reconstruction loss : {self.config.reconstruction_loss}")
         self.log(
-            f"Weighted sum : {self.config.n_cuts_weight}*Ncuts + {self.config.rec_loss_weight}*Reconstruction"
+            f"Weighted sum : {self.config.n_cuts_weight}*NCuts + {self.config.rec_loss_weight}*Reconstruction"
         )
         ##############
         self.log("-- Data --")
         self.log("Training data :")
         [
-            self.log(f"\n{v}")
+            self.log(f"{v}")
             for d in self.config.train_data_dict
             for k, v in d.items()
         ]
         if self.config.eval_volume_dict is not None:
             self.log("Validation data :")
             [
-                self.log(f"\n{k}: {v}")
+                self.log(f"{k}: {v}")
                 for d in self.config.eval_volume_dict
                 for k, v in d.items()
             ]
@@ -443,9 +441,9 @@ def train(self):
             set_track_meta(False)
             ##############
             # if WANDB_INSTALLED:
-            #     wandb.init(
-            #         config=WANDB_CONFIG, project="WNet-benchmark", mode=WANDB_MODE
-            #     )
+            # wandb.init(
+            #     config=WANDB_CONFIG, project="WNet-benchmark", mode=WANDB_MODE
+            # )
 
             set_determinism(
                 seed=self.config.deterministic_config.seed
@@ -455,12 +453,8 @@ def train(self):
             normalize_function = utils.remap_image
             device = self.config.device
 
-            self.log(f"Using device: {device}")
-
-            # self.log("Config:")  # FIXME log_parameters func instead
-            # [self.log(str(a)) for a in self.config.__dict__.items()]
+            # self.log(f"Using device: {device}")
             self.log_parameters()
-
             self.log("Initializing training...")
             self.log("Getting the data")
 
@@ -473,7 +467,6 @@ def train(self):
             #               Training the model                #
             ###################################################
             self.log("Initializing the model:")
-
             self.log("- Getting the model")
             # Initialize the model
             model = WNet(
@@ -494,8 +487,8 @@ def train(self):
                         )
                     )
 
-            if WANDB_INSTALLED:
-                wandb.watch(model, log_freq=100)
+            # if WANDB_INSTALLED:
+            #     wandb.watch(model, log_freq=100)
 
             if self.config.weights_info.custom:
                 if self.config.weights_info.use_pretrained:
@@ -619,10 +612,10 @@ def train(self):
                         )
 
                     epoch_rec_loss += reconstruction_loss.item()
-                    if WANDB_INSTALLED:
-                        wandb.log(
-                            {"Reconstruction loss": reconstruction_loss.item()}
-                        )
+                    # if WANDB_INSTALLED:
+                    #     wandb.log(
+                    #         {"Reconstruction loss": reconstruction_loss.item()}
+                    #     )
 
                     # Backward pass for the reconstruction loss
                     optimizer.zero_grad()
@@ -631,8 +624,8 @@ def train(self):
 
                     loss = alpha * Ncuts + beta * reconstruction_loss
                     epoch_loss += loss.item()
-                    if WANDB_INSTALLED:
-                        wandb.log({"Weighted sum of losses": loss.item()})
+                    # if WANDB_INSTALLED:
+                    #     wandb.log({"Weighted sum of losses": loss.item()})
                     loss.backward(loss)
                     optimizer.step()
 
@@ -818,9 +811,9 @@ def train(self):
                             self.log(f"Saving new best model to {save_path}")
                             torch.save(model.state_dict(), save_path)
 
-                        if WANDB_INSTALLED:
-                            # log validation dice score for each validation round
-                            wandb.log({"val/dice_metric": metric})
+                        # if WANDB_INSTALLED:
+                        # log validation dice score for each validation round
+                        # wandb.log({"val/dice_metric": metric})
 
                         dec_out_val = (
                             val_decoder_outputs[0].detach().cpu().numpy()

napari_cellseg3d/code_plugins/plugin_model_training.py

@@ -1402,7 +1402,7 @@ def _show_plot_max(self, plot, y):
             x_max,
             dice_max,
             c="r",
-            label="Max. Dice.",
+            label="Max. Dice",
             zorder=5,
         )