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().astype(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().astype(torch.float32).cpu().numpy()[0, 0]
+        masks_dict[view] = masks.detach().astype(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

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().astype(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().astype(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().astype(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().astype(torch.float32).cpu().numpy()))
         assert v.shape == (batch,)
 
     # ensure inputs are not modified