From 82682ceeb2fe10cb284ee3209abb6eb7fd5cf997 Mon Sep 17 00:00:00 2001
From: Laurent Fainsin <laurentfainsin@protonmail.com>
Date: Mon, 11 Jul 2022 17:02:13 +0200
Subject: [PATCH] feat: got precision 16 back

Former-commit-id: 6b19dc9bd17078bb2c151d5cd96e7ba4da9e1b89 [formerly 5d1eac2ed10be960c89407ad265ff350e11c1adf]
Former-commit-id: 1db4ca0ce11ac818408b94625b872c1202b5d4ed
---
 src/data/dataset.py   | 6 +++++-
 src/train.py          | 4 ++--
 src/unet/module.py    | 2 +-
 src/utils/callback.py | 2 +-
 src/utils/paste.py    | 4 ++--
 5 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/data/dataset.py b/src/data/dataset.py
index d30931f..61508f2 100644
--- a/src/data/dataset.py
+++ b/src/data/dataset.py
@@ -58,7 +58,7 @@ class LabeledDataset(Dataset):
 
         # open and convert mask
         mask_path = self.images[index].parent.joinpath("MASK.PNG")
-        mask = np.array(Image.open(mask_path).convert("L"), dtype=np.uint8) / 255
+        mask = np.array(Image.open(mask_path).convert("L"), dtype=np.uint8) // 255
 
         # convert image & mask to Tensor float in [0, 1]
         post_process = A.Compose(
@@ -72,4 +72,8 @@ class LabeledDataset(Dataset):
         image = augmentations["image"]
         mask = augmentations["mask"]
 
+        # make sure image and mask are floats, TODO: mettre dans le post_process, ToFloat Image only
+        image = image.float()
+        mask = mask.float()
+
         return image, mask
diff --git a/src/train.py b/src/train.py
index 85625e0..7167634 100644
--- a/src/train.py
+++ b/src/train.py
@@ -38,7 +38,7 @@ if __name__ == "__main__":
     # model.load_state_dict(state_dict)
 
     # log gradients and weights regularly
-    logger.watch(model, log="all")
+    logger.watch(model.model, log="all")
 
     # Create the dataloaders
     datamodule = Spheres()
@@ -49,7 +49,7 @@ if __name__ == "__main__":
         accelerator=wandb.config.DEVICE,
         benchmark=wandb.config.BENCHMARK,
         # profiler="simple",
-        # precision=16,
+        precision=16,
         logger=logger,
         log_every_n_steps=1,
         val_check_interval=100,
diff --git a/src/unet/module.py b/src/unet/module.py
index c7ae528..43b09d0 100644
--- a/src/unet/module.py
+++ b/src/unet/module.py
@@ -38,7 +38,7 @@ class UNetModule(pl.LightningModule):
 
         # forward pass, compute masks
         prediction = self.model(data)
-        binary = (torch.sigmoid(prediction) > 0.5).float()  # TODO: check if float necessary
+        binary = (torch.sigmoid(prediction) > 0.5).half()
 
         # compute metrics (in dictionnary)
         metrics = {
diff --git a/src/utils/callback.py b/src/utils/callback.py
index 20497f2..969b619 100644
--- a/src/utils/callback.py
+++ b/src/utils/callback.py
@@ -31,7 +31,7 @@ class TableLog(Callback):
                 zip(
                     images.cpu(),
                     ground_truth.cpu(),
-                    predictions["linear"].cpu(),
+                    predictions["linear"].cpu().float(),
                     predictions["binary"].cpu().squeeze(1).int().numpy(),
                 )
             ):
diff --git a/src/utils/paste.py b/src/utils/paste.py
index 02c1a1b..be6f25d 100644
--- a/src/utils/paste.py
+++ b/src/utils/paste.py
@@ -87,7 +87,7 @@ class RandomPaste(A.DualTransform):
 
             img.paste(paste_img, (x, y), paste_mask)
 
-        return np.asarray(img.convert("RGB"))
+        return np.array(img.convert("RGB"))
 
     def apply_to_mask(self, mask, augmentations, paste_mask, **params):
         # convert mask to Image, needed for `paste` function
@@ -116,7 +116,7 @@ class RandomPaste(A.DualTransform):
 
             mask.paste(paste_mask, (x, y), paste_mask_bin)
 
-        return np.asarray(mask.convert("L"))
+        return np.array(mask.convert("L"))
 
     def get_params_dependent_on_targets(self, params):
         # choose a random image and its corresponding mask