minor changes

Author: nauyan

config.py

@@ -14,7 +14,8 @@
 checkpoints_dir = "checkpoints"
 tensorboard_logs = "run_logs"
 
-encoder = "tu-tf_efficientnet_b5_ns"
+# encoder = "tu-tf_efficientnet_b5_ns"
+encoder = "tu-tf_efficientnet_b0_ns"
 encoder_weights = None
 
 # CoNSep Weights
@@ -24,7 +25,7 @@
 inference_weights = "NC-Net_all_metric.pth"
 
 device = "cuda"
-batch_size = 16
+batch_size = 32
 epochs = 1000
 learning_rate = 0.01
 watershed_threshold = 0.2

src/metrics.py

@@ -11,22 +11,24 @@ def dice_score(pred, true):
     pred = F.softmax(pred.permute(0, 2, 3, 1).contiguous(), dim=-1)
     pred = torch.argmax(pred, dim=-1, keepdim=False)
 
-    true = np.copy(true.cpu().detach().numpy())
-    pred = np.copy(pred.cpu().detach().numpy())
+    # true = np.copy(true.cpu().detach().numpy())
+    # pred = np.copy(pred.cpu().detach().numpy())
 
     inter = true * pred
     denom = true + pred
-    return (2.0 * np.sum(inter)) / np.sum(denom)
+
+    score = (2.0 * torch.sum(inter)) / torch.sum(denom)
+    return score.cpu().detach().numpy()
 
 
 def mse_metric(pred, true):
 
     pred = pred[:, 2, :, :]
     true = true[:, :, :, 1]
 
-    true = np.copy(true.cpu().detach().numpy())
-    pred = np.copy(pred.cpu().detach().numpy())
+    # true = np.copy(true.cpu().detach().numpy())
+    # pred = np.copy(pred.cpu().detach().numpy())
 
     loss = pred - true
     loss = (loss * loss).mean()
-    return loss
+    return loss.cpu().detach().numpy()

train.py

@@ -44,11 +44,11 @@ def train(train_dir, test_dir, dataset_name):
     train_loader = DataLoader(
         train_dataset,
         batch_size=batch_size,
-        shuffle=True,
+        shuffle=False,
         num_workers=multiprocessing.cpu_count(),
         pin_memory=True,
         persistent_workers=True,
-        prefetch_factor=8,
+        prefetch_factor=32,
     )
     valid_loader = DataLoader(
         valid_dataset,
@@ -88,8 +88,9 @@ def train(train_dir, test_dir, dataset_name):
         verbose=True,
     )
 
-    writer_path = "./{}/NC-Net_{}".format(config.tensorboard_logs,
-                                          dataset_name)
+    writer_path = "./{}/NC-Net_{}_{}".format(config.tensorboard_logs,
+                                            config.encoder,
+                                            dataset_name)
     writer = SummaryWriter(writer_path)
     min_loss = 9999
     max_score = 0
@@ -110,7 +111,8 @@ def train(train_dir, test_dir, dataset_name):
             min_loss = valid_logs[loss_fn.__name__]
             torch.save(
                 model.state_dict(),
-                "./{}/NC-Net_{}.pth".format(config.checkpoints_dir,
+                "./{}/NC-Net_{}_{}.pth".format(config.checkpoints_dir,
+                                            config.encoder,
                                             dataset_name),
             )
             last_save = i
@@ -120,8 +122,9 @@ def train(train_dir, test_dir, dataset_name):
             max_score = valid_logs[metrics[0].__name__]
             torch.save(
                 model.state_dict(),
-                "./{}/NC-Net_{}_metric.pth".format(config.checkpoints_dir,
-                                                   dataset_name),
+                "./{}/NC-Net__{}_{}_metric.pth".format(config.checkpoints_dir,
+                                                    config.encoder,
+                                                    dataset_name),
             )
             last_save = i
             print("Model saved Metric!")