Fix dtype

Author: mathpluscode

cinema/examples/inference/landmark_heatmap.py

@@ -39,7 +39,7 @@ def run(view: str, seed: int, device: torch.device, dtype: torch.dtype) -> None:
         with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
             logits = model(batch)[view]  # (1, 3, x, y)
         probs = torch.sigmoid(logits)  # (1, 3, width, height)
-        probs_list.append(probs[0].detach().cpu().numpy())
+        probs_list.append(probs[0].detach().to(torch.float32).cpu().numpy())
         coords = heatmap_soft_argmax(probs)[0].numpy()
         coords = [int(x) for x in coords]
 

cinema/examples/inference/mae.py

@@ -77,8 +77,8 @@ def reconstruct_images(
             patch_size=patch_size_dict[view],
             grid_size=grid_size_dict[view],
         )
-        reconstructed_dict[view] = reconstructed.detach().cpu().numpy()[0, 0]
-        masks_dict[view] = masks.detach().cpu().numpy()[0, 0]
+        reconstructed_dict[view] = reconstructed.detach().to(torch.float32).cpu().numpy()[0, 0]
+        masks_dict[view] = masks.detach().to(torch.float32).cpu().numpy()[0, 0]
     reconstructed_dict["sax"] = reconstructed_dict["sax"][..., :sax_slices]
     masks_dict["sax"] = masks_dict["sax"][..., :sax_slices]
     return reconstructed_dict, masks_dict
@@ -115,12 +115,12 @@ def run(device: torch.device, dtype: torch.dtype) -> None:
     lax_4c_image = np.transpose(sitk.GetArrayFromImage(sitk.ReadImage(exp_dir / "data/ukb/1/1_lax_4c.nii.gz")))
 
     image_dict = {
-        "sax": sax_image[..., t],
-        "lax_2c": lax_2c_image[..., 0, t],
-        "lax_3c": lax_3c_image[..., 0, t],
-        "lax_4c": lax_4c_image[..., 0, t],
+        "sax": sax_image[None, ..., t],
+        "lax_2c": lax_2c_image[None, ..., 0, t],
+        "lax_3c": lax_3c_image[None, ..., 0, t],
+        "lax_4c": lax_4c_image[None, ..., 0, t],
     }
-    batch = {k: torch.from_numpy(v[None, ...]) for k, v in image_dict.items()}
+    batch = {k: torch.from_numpy(v) for k, v in image_dict.items()}
 
     # forward
     sax_slices = batch["sax"].shape[-1]
@@ -137,10 +137,12 @@ def run(device: torch.device, dtype: torch.dtype) -> None:
             grid_size_dict,
             sax_slices,
         )
+        batch = {k: v.detach().to(torch.float32).cpu().numpy()[0, 0] for k, v in batch.items()}
+        batch["sax"] = batch["sax"][..., :sax_slices]
 
     # visualize
     fig = plot_mae_reconstruction(
-        image_dict,
+        batch,
         reconstructed_dict,
         masks_dict,
     )

cinema/examples/inference/segmentation_lax_4c.py

@@ -126,7 +126,7 @@ def run(trained_dataset: str, seed: int, device: torch.device, dtype: torch.dtyp
         batch = {k: v[None, ...].to(device=device, dtype=dtype) for k, v in batch.items()}
         with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
             logits = model(batch)[view]  # (1, 4, x, y)
-        labels = torch.argmax(logits, dim=1)[0].detach().cpu().numpy()  # (x, y)
+        labels = torch.argmax(logits, dim=1)[0].detach().to(torch.float32).cpu().numpy()  # (x, y)
 
         # the model seems to hallucinate an additional right ventricle and myocardium sometimes
         # find the connected component that is closest to left ventricle

cinema/examples/inference/segmentation_sax.py

@@ -107,7 +107,7 @@ def run(trained_dataset: str, seed: int, device: torch.device, dtype: torch.dtyp
         with torch.no_grad(), torch.autocast("cuda", dtype=dtype, enabled=torch.cuda.is_available()):
             logits = model(batch)[view]  # (1, 4, x, y, z)
         labels_list.append(torch.argmax(logits, dim=1)[0, ..., :n_slices])
-    labels = torch.stack(labels_list, dim=-1).detach().cpu().numpy()  # (x, y, z, t)
+    labels = torch.stack(labels_list, dim=-1).detach().to(torch.float32).cpu().numpy()  # (x, y, z, t)
 
     # visualise segmentations
     fig = plot_segmentations(images, labels, t_step)

cinema/mae/mae_test.py

@@ -125,7 +125,7 @@ def test_conv_mae_size(
     # value can be nan if target is empty
     # this is unlikely to happen with large mask_ratio
     if min(ns_masked) > 0:
-        assert not np.isnan(loss.detach().cpu().numpy())
+        assert not np.isnan(loss.detach().to(torch.float32).cpu().numpy())
         for v in metrics.values():
             assert not np.isnan(v.detach())
             assert v.shape == ()

cinema/segmentation/train_test.py

@@ -105,7 +105,7 @@ def test_segmentation_eval_metrics(
 
     metrics = segmentation_metrics(logits, labels, spacing)
     for v in metrics.values():
-        assert not np.any(np.isnan(v.detach().cpu().numpy()))
+        assert not np.any(np.isnan(v.detach().to(torch.float32).cpu().numpy()))
         assert v.shape == (batch,)
 
     # ensure inputs are not modified