add validation tests

dacapo/predict_local.py

@@ -71,6 +71,10 @@ def predict(
     compute_context = create_compute_context()
     device = compute_context.device
 
+    model_device = next(model.parameters()).device
+
+    assert model_device == device, f"Model is not on the right device, Model: {model_device}, Compute device: {device}" 
+
     def predict_fn(block):
         raw_input = raw_array.to_ndarray(block.read_roi)
 
@@ -85,6 +89,7 @@ def predict_fn(block):
             raw_input = np.expand_dims(raw_input, 0)
             axis_names = ["c^"] + axis_names
 
+
         with torch.no_grad():
             model.eval()
             predictions = (

tests/operations/test_architecture.py

@@ -2,6 +2,7 @@
 
 import pytest
 from pytest_lazy_fixtures import lf
+import torch.nn as nn
 
 import logging
 
@@ -27,6 +28,7 @@ def test_architecture(
     [
         lf("dummy_architecture"),
         lf("unet_architecture"),
+        lf("unet_3d_architecture"),
     ],
 )
 def test_stored_architecture(
@@ -48,3 +50,33 @@ def test_stored_architecture(
     architecture = retrieved_arch_config.architecture_type(retrieved_arch_config)
 
     assert architecture.dims is not None, f"Architecture dims are None {architecture}"
+
+
+@pytest.mark.parametrize(
+    "architecture_config",
+    [
+        lf("unet_architecture"),
+    ],
+)
+def test_2d_conv_unet(
+    architecture_config,
+):
+    architecture = architecture_config.architecture_type(architecture_config)
+    for name, module in architecture.named_modules():
+        if isinstance(module, nn.Conv3d):
+            raise ValueError(f"Conv3d found in 2d unet {name}")
+
+
+@pytest.mark.parametrize(
+    "architecture_config",
+    [
+        lf("unet_3d_architecture"),
+    ],
+)
+def test_2d_conv_unet(
+    architecture_config,
+):
+    architecture = architecture_config.architecture_type(architecture_config)
+    for name, module in architecture.named_modules():
+        if isinstance(module, nn.Conv2d):
+            raise ValueError(f"Conv2d found in 3d unet {name}")

tests/operations/test_validate.py

@@ -130,4 +130,4 @@ def test_validate_unet(datasplit, task, trainer, architecture):
 
     # -------------------------------------
     weights_store.store_weights(run, 0)
-    validate_run(run, 0)
+    validate(run.name, 0)