Removed softmax-layer.output from trainer.

deepmedic/frontEnd/trainSession.py

@@ -121,6 +121,7 @@ def run_session(self, *args):
             with tf.compat.v1.variable_scope("trainer"):
                 self._log.print3("=========== Building Trainer ===========\n")
                 trainer = Trainer(*(self._params.get_args_for_trainer() + [cnn3d]))
+                trainer.compute_costs(self._log, p_y_given_x_train)
                 trainer.create_optimizer(*self._params.get_args_for_optimizer())  # Trainer and net connect here.
 
             tensorboard_loggers = self.create_tensorboard_loggers(['train', 'val'],

deepmedic/neuralnet/cnn3d.py

@@ -401,9 +401,6 @@ def apply(self, inputs_per_pathw, mode, train_val_test, verbose=True, log=None):
         # Softmax
         p_y_given_x = self.finalTargetLayer.apply(logits_no_bias, mode)
 
-        # TODO: REMOVE THE BELOW for eager mode. NEEDED by the trainer.py (loss) for now. And returning FMs.
-        self.finalTargetLayer.output[train_val_test] = p_y_given_x
-
         return p_y_given_x
 
     def calc_inp_dims_of_paths_from_hr_inp(self, inp_hr_dims):

deepmedic/neuralnet/trainer.py

@@ -26,17 +26,15 @@ class Trainer(object):
     # This acts as a learning-rate scheduler and as a manager of the optimizer.
     # In case of multiple optimizers, either this must become a manager of multiple optimizers+schedules, ...
     # ... or I ll need one trainer per optimizer and a new trainer-manager. 
-    def __init__(   self,
-                    log,
-                    indicesOfLayersPerPathwayTypeToFreeze,
-                    losses_and_weights,
-                    L1_reg_weight,
-                    L2_reg_weight,
-                    # Cost schedules
-                    reweight_classes_in_cost,
-
-                    network_to_train
-                    ):
+    def __init__(self,
+                 log,
+                 indicesOfLayersPerPathwayTypeToFreeze,
+                 losses_and_weights,
+                 L1_reg_weight,
+                 L2_reg_weight,
+                 # Cost schedules
+                 reweight_classes_in_cost,
+                 network_to_train):
 
         log.print3("Building Trainer.")
 
@@ -95,12 +93,9 @@ def __init__(   self,
         self._op_assign_epoch_with_top_mean_val_acc_tvf = None
         self._op_assign_last_epoch_lr_lowered = None
 
-        # Setup any parameters needed
-        self._setup_costs(log) # Needs to be run with initialized self._num_epochs_trained_tfv
-
 
     ############## All the logic wrt cost / regularizers should be done here ##############
-    def _setup_costs(self, log):
+    def compute_costs(self, log, p_y_given_x): # Needs to be run with initialized self._num_epochs_trained_tfv
         if not self._total_cost is None:
             log.print3("ERROR: Problem in Trainer. It was called to setup the total cost, but it was not None."+\
                        "\n\t This should not happen. Setup should be called only once.\n Exiting!")
@@ -111,14 +106,14 @@ def _setup_costs(self, log):
         y_gt = self._net._output_gt_tensor_feeds['train']['y_gt']
         if "xentr" in self._losses_and_weights and self._losses_and_weights["xentr"] is not None:
             log.print3("COST: Using cross entropy with weight: " +str(self._losses_and_weights["xentr"]))
-            w_per_cl_vec = self._compute_w_per_class_vector_for_xentr( self._net.num_classes, y_gt )
-            cost += self._losses_and_weights["xentr"] * cfs.x_entr( self._net.finalTargetLayer.output["train"], y_gt, w_per_cl_vec )
+            w_per_cl_vec = self._compute_w_per_class_vector_for_xentr(self._net.num_classes, y_gt)
+            cost += self._losses_and_weights["xentr"] * cfs.x_entr(p_y_given_x, y_gt, w_per_cl_vec)
         if "iou" in self._losses_and_weights and self._losses_and_weights["iou"] is not None:
             log.print3("COST: Using iou loss with weight: " +str(self._losses_and_weights["iou"]))
-            cost += self._losses_and_weights["iou"] * cfs.iou( self._net.finalTargetLayer.output["train"], y_gt )
+            cost += self._losses_and_weights["iou"] * cfs.iou(p_y_given_x, y_gt)
         if "dsc" in self._losses_and_weights and self._losses_and_weights["dsc"] is not None:
             log.print3("COST: Using dsc loss with weight: " +str(self._losses_and_weights["dsc"]))
-            cost += self._losses_and_weights["dsc"] * cfs.dsc( self._net.finalTargetLayer.output["train"], y_gt )
+            cost += self._losses_and_weights["dsc"] * cfs.dsc(p_y_given_x, y_gt)
 
         cost_L1_reg = self._L1_reg_weight * cfs.cost_L1(self._net.params_for_L1_L2_reg())
         cost_L2_reg = self._L2_reg_weight * cfs.cost_L2(self._net.params_for_L1_L2_reg())