tests: keep python tests alive, they're still useful sometimes

Author: Acly

tests/test_birefnet.py

@@ -228,10 +228,11 @@ def __init__(
             drop=drop,
         )
 
-        self.H = None
-        self.W = None
+        self.H: int | None = None
+        self.W: int | None = None
 
     def forward(self, x, mask_matrix):
+        assert self.W is not None and self.H is not None, "W and H must be set before forward"
         B, L, C = x.shape
         H, W = self.H, self.W
         assert L == H * W, "input feature has wrong size"
@@ -297,7 +298,7 @@ def test_swin_block():
 
     x = input_tensor(1, 36, 8)
     mask = torch.zeros(2, 9, 9).masked_fill(torch.rand(2, 9, 9) > 0.5, -100.0)
-    state["mask"] = mask
+    state["mask"] = mask.half()
     swin_block.W, swin_block.H = 6, 6
     expected = swin_block(x, None)
 
@@ -421,7 +422,7 @@ def attention_mask(self, H, W):
         mask_windows = window_partition(img_mask, self.window_size)
         mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
         attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
-        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0))
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float("-inf"))
         attn_mask = attn_mask.masked_fill(attn_mask == 0, float(0.0))
         return attn_mask
 
@@ -453,7 +454,7 @@ def forward(self, x, H, W):
         mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
         attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
         attn_mask = (
-            attn_mask.masked_fill(attn_mask != 0, float(-100.0))
+            attn_mask.masked_fill(attn_mask != 0, float("-inf"))
             .masked_fill(attn_mask == 0, float(0.0))
             .to(x.dtype)
         )
@@ -475,8 +476,8 @@ def test_attention_mask():
     swin_layer = BasicLayer(8, 2, 2, window_size=window_size)
     expected = swin_layer.attention_mask(h, w)
 
-    result = torch.zeros_like(expected)
-    result = workbench.invoke_test("biref_attention_mask", result, {})
+    x = torch.zeros_like(expected)
+    result = workbench.invoke_test("biref_attention_mask", x, {})
 
     assert torch.allclose(result, expected)
 

tests/test_mobile_sam.py

@@ -1325,7 +1325,8 @@ def test_output_upscaling():
     result = workbench.invoke_test("sam_output_upscaling", x, state, nhwc_layout, backend="vulkan")
     result = to_nchw(result)
 
-    assert torch.allclose(result, expected, atol=1e-4, rtol=1e-2)  # fp16 weights
+    workbench.print_results(result, expected)
+    assert torch.allclose(result, expected, rtol=0.1)  # fp16 weights
 
 
 class MaskDecoder(torch.nn.Module):

tests/test_primitives.py

@@ -2,7 +2,7 @@
 import torch
 
 from . import workbench
-from .workbench import to_nchw, to_nhwc
+from .workbench import input_tensor, to_nchw, to_nhwc
 
 
 def test_linear():
@@ -43,56 +43,59 @@ def test_conv_2d_depthwise(scenario: str, memory_layout: str, batch: str, backen
         x = to_nhwc(x)
         k = k.permute(2, 3, 1, 0)
     test_case = f"conv_2d_depthwise_{memory_layout}"
-    params = dict(stride=stride, pad=pad, dilation=dilate)
+    params = dict(stride=stride, pad=pad, dilation=dilate, memory_layout=memory_layout)
     result = workbench.invoke_test(test_case, x, dict(weight=k), params, backend)
     if memory_layout == "nhwc":
         result = to_nchw(result)
 
     assert torch.allclose(result, expected)
 
 
-@pytest.mark.parametrize("scenario", ["3x3", "5x5", "stride2"])
+@pytest.mark.parametrize("scenario", ["3x3", "5x5", "stride2", "nhwc"])
 def test_conv_transpose_2d(scenario: str):
     ksize, stride = {
         "3x3": (3, 1),
         "5x5": (5, 1),
         "stride2": (3, 2),
-        "nchw": (3, 1),
+        "nhwc": (3, 1),
     }[scenario]
-    x = torch.arange(2 * 11 * 4 * 5).reshape(2, 11, 4, 5).float()
-    weight = torch.arange(11 * 2 * ksize * ksize).reshape(11, 2, ksize, ksize).float()
+    x = input_tensor(2, 11, 4, 5)
+    weight = input_tensor(11, 2, ksize, ksize)
     bias = None
     expected = torch.nn.functional.conv_transpose2d(x, weight, bias, stride=stride)
 
-    x = to_nhwc(x)  # -> [N, H, W, C_in]
+    if scenario == "nhwc":
+        x = to_nhwc(x)  # -> [N, H, W, C_in]
     result = workbench.invoke_test(
         "conv_transpose_2d",
         x,
         dict(weight=weight),
-        dict(stride=stride),
+        dict(stride=stride, memory_layout="nhwc" if scenario == "nhwc" else "nchw"),
         backend="vulkan",
     )
-    result = to_nchw(result)
+    if scenario == "nhwc":
+        result = to_nchw(result)
 
-    assert torch.allclose(result, expected)
+    workbench.print_results(result, expected)
+    assert torch.allclose(result, expected, rtol=1e-2)
 
 
-def test_batch_norm_2d():
-    x = torch.rand(1, 3, 4, 5)
-    weight = torch.rand(3)
-    bias = torch.rand(3)
-    mean = torch.rand(3)
-    var = torch.arange(1, 4).float()
-    expected = torch.nn.functional.batch_norm(x, mean, var, weight, bias, eps=1e-5)
+# def test_batch_norm_2d():
+#     x = torch.rand(1, 3, 4, 5)
+#     weight = torch.rand(3)
+#     bias = torch.rand(3)
+#     mean = torch.rand(3)
+#     var = torch.arange(1, 4).float()
+#     expected = torch.nn.functional.batch_norm(x, mean, var, weight, bias, eps=1e-5)
 
-    x = to_nhwc(x)
+#     x = to_nhwc(x)
 
-    var = (var + 1e-5).sqrt()
-    state = dict(weight=weight, bias=bias, running_mean=mean, running_var=var)
-    result = workbench.invoke_test("batch_norm_2d", x, state)
-    result = to_nchw(result)
+#     var = (var + 1e-5).sqrt()
+#     state = dict(weight=weight, bias=bias, running_mean=mean, running_var=var)
+#     result = workbench.invoke_test("batch_norm_2d", x, state, dict(memory_layout="nhwc"))
+#     result = to_nchw(result)
 
-    assert torch.allclose(result, expected)
+#     assert torch.allclose(result, expected)
 
 
 def test_layer_norm():

tests/workbench.cpp

@@ -284,10 +284,11 @@ DEF(biref_patch_merging)(model_ref m, span<tensor> input, param_dict const& p) {
     return {swin::patch_merging(m, input[0], 6, 4)};
 }
 
-DEF(biref_attention_mask)(model_ref m, span<tensor> input, param_dict const& p) {
-    auto dst = span((byte*)input[0]->data, ggml_nbytes(input[0]));
-    swin::compute_attention_mask(dst, 18, 18, 6);
-    return {input[0]};
+DEF(biref_attention_mask)(model_ref m, span<tensor> /*input*/, param_dict const& p) {
+    auto mask = swin::create_attention_mask(m, 18, 18, 6);
+    ggml_backend_alloc_ctx_tensors(m, workbench_backend());
+    transfer_to_backend(mask);
+    return {ggml_cast(m, mask.x, GGML_TYPE_F32)};
 }
 
 DEF(biref_swin_layer)(model_ref m, span<tensor> input, param_dict const& p) {

tests/workbench.py

@@ -1,5 +1,6 @@
 import ctypes
 from functools import reduce
+from typing import Mapping
 import torch
 import os
 
@@ -66,7 +67,7 @@ def raw_to_torch_tensor(raw_tensor: RawTensor):
     ).reshape(shape)
 
 
-def encode_params(params: dict[str, str | int | float]):
+def encode_params(params: Mapping[str, str | int | float]):
     raw_params = []
     for name, value in params.items():
         ptype = 0
@@ -109,7 +110,7 @@ def invoke_test(
     test_case: str,
     input: torch.Tensor | list[torch.Tensor],
     state: dict[str, torch.Tensor],
-    params: dict[str, str | int | float] = {},
+    params: Mapping[str, str | int | float] = {},
     backend: str = "cpu",
 ):
     input = input if isinstance(input, list) else [input]
@@ -142,7 +143,7 @@ def invoke_test(
     return output
 
 
-def input_tensor(*shape: tuple[int]):
+def input_tensor(*shape: int):
     end = reduce(lambda x, y: x * y, shape, 1)
     return torch.arange(0, end).reshape(*shape) / end