interpolation arg

Author: Mayukhdeb

tests/test_laplacian_pyramid.py

@@ -15,14 +15,15 @@
 @pytest.mark.parametrize("dtype2", supported_dtypes)
 @pytest.mark.parametrize("height, width", [(16, 16), (32, 32)])
 @pytest.mark.parametrize("factor_w, factor_h", [(2.0, 2.0), (3.0, 3.0), (4.0, 4.0)])
-def test_laplacian_pyramid_loss_precision(dtype1, dtype2, height, width, factor_w, factor_h):
+@pytest.mark.parametrize("interpolation", ["bilinear", "nearest"])
+def test_laplacian_pyramid_loss_precision(dtype1, dtype2, height, width, factor_w, factor_h, interpolation):
     # Create a sample cortical_sheet tensor
     e = 16  # Example depth
     torch.manual_seed(42)  # Set seed for reproducibility
     cortical_sheet = torch.rand(height, width, e).to(dtype1)
 
     # Call the function with the given precision
-    loss = laplacian_pyramid_loss(cortical_sheet, factor_w=factor_w, factor_h=factor_h)
+    loss = laplacian_pyramid_loss(cortical_sheet, factor_w=factor_w, factor_h=factor_h, interpolation=interpolation)
 
     # Ensure the loss is finite
     assert torch.isfinite(loss), f"Loss is not finite for dtype {dtype1}"
@@ -33,7 +34,7 @@ def test_laplacian_pyramid_loss_precision(dtype1, dtype2, height, width, factor_
     # Compare results to float32 (considered ground truth for higher precision)
     if dtype1 !=dtype2:
         float32_sheet = cortical_sheet.to(dtype2)
-        expected_loss = laplacian_pyramid_loss(cortical_sheet=float32_sheet, factor_w=factor_w, factor_h=factor_h)
+        expected_loss = laplacian_pyramid_loss(cortical_sheet=float32_sheet, factor_w=factor_w, factor_h=factor_h, interpolation=interpolation)
         assert_close(
             loss.to(dtype2), expected_loss, rtol=1e-3, atol=1e-4,
             msg=f"Loss mismatch for dtype1 {dtype1} loss: {loss.to(dtype2)} and dtype2 {dtype2} loss2: {expected_loss}"

tests/test_loss_conv.py

@@ -15,8 +15,9 @@
 @pytest.mark.parametrize("hidden_channels", [16, 32])
 @pytest.mark.parametrize("init_from_layer", [True, False])
 @pytest.mark.parametrize("dtype", supported_dtypes)
+@pytest.mark.parametrize("interpolation", ["bilinear", "nearest"])
 def test_loss_conv(
-    num_steps: int, hidden_channels: int, init_from_layer: bool, dtype
+    num_steps: int, hidden_channels: int, init_from_layer: bool, dtype, interpolation: str
 ):  # num_steps is now passed by the fixture
 
     # Define the model
@@ -30,16 +31,16 @@ def test_loss_conv(
     if init_from_layer:
         losses = [
             LaplacianPyramid.from_layer(
-                model=model, layer=model[0], scale=1.0, factor_h=3.0, factor_w=3.0
+                model=model, layer=model[0], scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation
             ),
             LaplacianPyramid.from_layer(
-                model=model, layer=model[2], scale=1.0, factor_h=3.0, factor_w=3.0
+                model=model, layer=model[2], scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation
             ),
         ]
     else:
         losses = [
-            LaplacianPyramid(layer_name="0", scale=1.0, factor_h=3.0, factor_w=3.0),
-            LaplacianPyramid(layer_name="2", scale=1.0, factor_h=3.0, factor_w=3.0),
+            LaplacianPyramid(layer_name="0", scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation),
+            LaplacianPyramid(layer_name="2", scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation),
         ]
 
     # Define the TopoLoss

topoloss/core.py

@@ -44,6 +44,7 @@ def get_layerwise_topo_losses(self, model, do_scaling: bool = True) -> dict:
                     cortical_sheet=cortical_sheet,
                     factor_h=loss_info.factor_h,
                     factor_w=loss_info.factor_w,
+                    interpolation=loss_info.interpolation,
                 )
             elif isinstance(loss_info, LaplacianPyramidOnBias):
                 assert isinstance(
@@ -61,6 +62,7 @@ def get_layerwise_topo_losses(self, model, do_scaling: bool = True) -> dict:
                     cortical_sheet=cortical_sheet,
                     factor_h=loss_info.factor_h,
                     factor_w=loss_info.factor_w,
+                    interpolation=loss_info.interpolation,
                 )
 
             elif isinstance(loss_info, LaplacianPyramidOnInput):
@@ -72,6 +74,7 @@ def get_layerwise_topo_losses(self, model, do_scaling: bool = True) -> dict:
                     cortical_sheet=cortical_sheet,
                     factor_h=loss_info.factor_h,
                     factor_w=loss_info.factor_w,
+                    interpolation=loss_info.interpolation,
                 )
 
             if do_scaling:

topoloss/losses/laplacian_pyramid.py

@@ -10,7 +10,7 @@
 import torch
 
 def laplacian_pyramid_loss(
-    cortical_sheet: TensorType["height", "width", "e"], factor_w: float, factor_h: float
+    cortical_sheet: TensorType["height", "width", "e"], factor_w: float, factor_h: float, interpolation: str = "bilinear"
 ):
     grid = cortical_sheet
     assert grid.ndim == 3, "Expected grid to be a 3d tensor of shape (h, w, e)"
@@ -24,10 +24,10 @@ def laplacian_pyramid_loss(
     ), f"Expected factor_w to be <= grid.shape[2] = {grid.shape[2]} but got: {factor_w}"
     # Downscale the grid tensor
     downscaled_grid = F.interpolate(
-        grid, scale_factor=(1 / factor_h, 1 / factor_w), mode="bilinear"
+        grid, scale_factor=(1 / factor_h, 1 / factor_w), mode=interpolation
     )
     # Upscale the downscaled grid tensor
-    upscaled_grid = F.interpolate(downscaled_grid, size=grid.shape[2:], mode="bilinear")
+    upscaled_grid = F.interpolate(downscaled_grid, size=grid.shape[2:], mode=interpolation)
 
     # Calculate the MSE loss between the original grid and upscaled grid
     # loss = F.mse_loss(upscaled_grid, grid)
@@ -50,21 +50,23 @@ class LaplacianPyramid:
     layer_name: str
     factor_h: float
     factor_w: float
+    interpolation: str = "bilinear"
     scale: Optional[Union[None, float]] = field(default=1.0)
 
     @classmethod
-    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0):
+    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation: str ="bilinear"):
         layer_name = get_name_by_layer(model=model, layer=layer)
         return cls(
             layer_name=layer_name,
             scale=scale,
             factor_h=factor_h,
             factor_w=factor_w,
+            interpolation=interpolation
         )
 
 
 def laplacian_pyramid_loss_on_bias(
-    cortical_sheet: TensorType["h", "w"], factor_w: float, factor_h: float
+    cortical_sheet: TensorType["h", "w"], factor_w: float, factor_h: float, interpolation: str = "bilinear"
 ):
 
     grid = cortical_sheet
@@ -81,10 +83,10 @@ def laplacian_pyramid_loss_on_bias(
     grid = grid.unsqueeze(0).unsqueeze(0)
     # Downscale the grid tensor
     downscaled_grid = F.interpolate(
-        grid, scale_factor=(1 / factor_h, 1 / factor_w), mode="bilinear"
+        grid, scale_factor=(1 / factor_h, 1 / factor_w), mode=interpolation
     )
     # Upscale the downscaled grid tensor
-    upscaled_grid = F.interpolate(downscaled_grid, size=grid.shape[2:], mode="bilinear")
+    upscaled_grid = F.interpolate(downscaled_grid, size=grid.shape[2:], mode=interpolation)
 
     grid = rearrange(grid.squeeze(0).squeeze(0), "h w -> (h w)").unsqueeze(0)
     upscaled_grid = rearrange(
@@ -107,10 +109,11 @@ class LaplacianPyramidOnBias:
     layer_name: str
     factor_h: float
     factor_w: float
+    interpolation: str = "bilinear"
     scale: Optional[Union[None, float]] = field(default=1.0)
 
     @classmethod
-    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0):
+    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation: str ="bilinear"):
         assert (
             layer.bias is not None
         ), "Expected layer to have a bias, but got None. *sad sad sad*"
@@ -120,6 +123,7 @@ def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0):
             scale=scale,
             factor_h=factor_h,
             factor_w=factor_w,
+            interpolation=interpolation
         )
 
 @dataclass
@@ -136,10 +140,11 @@ class LaplacianPyramidOnInput:
     layer_name: str
     factor_h: float
     factor_w: float
+    interpolation: str = "bilinear"
     scale: Optional[Union[None, float]] = field(default=1.0)
 
     @classmethod
-    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0):
+    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation: str ="bilinear"):
         assert (
             layer.bias is not None
         ), "Expected layer to have a bias, but got None. *sad sad sad*"
@@ -149,4 +154,5 @@ def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0):
             scale=scale,
             factor_h=factor_h,
             factor_w=factor_w,
+            interpolation=interpolation
         )