File size: 4,601 Bytes
45d3d29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | #!/usr/bin/env python3
"""CPU regression tests for the paper mean-only color-match float-domain shift."""
from __future__ import annotations
import sys
import unittest
from pathlib import Path
import torch
DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))
from benchmark.verify_paper_reproduction import (
MEAN_ONLY_PROTOCOL as VERIFIER_MEAN_ONLY_PROTOCOL,
)
from scripts.infer import (
COLOR_MATCH_FORMULA,
MEAN_ONLY_PROTOCOL,
MODEL_CHOICES,
match_color_mean,
)
class MatchColorMeanTests(unittest.TestCase):
def test_shifts_each_channel_by_its_own_mean_difference(self) -> None:
sr = torch.tensor(
[[[[0.25, 0.5], [0.75, 1.0]], [[0.125, 0.25], [0.5, 0.875]], [[0.0, 0.5], [0.25, 0.5]]]],
dtype=torch.float32,
)
lr = torch.tensor(
[[[[0.75, 0.75], [0.75, 0.75]], [[0.25, 0.25], [0.25, 0.25]], [[0.5, 0.5], [0.5, 0.5]]]],
dtype=torch.float32,
)
shifted = match_color_mean(sr, lr)
expected = torch.tensor(
[[[[0.375, 0.625], [0.875, 1.0]], [[0.0, 0.0625], [0.3125, 0.6875]], [[0.1875, 0.6875], [0.4375, 0.6875]]]],
dtype=torch.float32,
)
torch.testing.assert_close(shifted, expected, rtol=0.0, atol=0.0)
def test_zero_shift_leaves_values_untouched(self) -> None:
sr = torch.tensor([[[[0.25, 0.5], [0.5, 0.75]]]], dtype=torch.float32)
lr = torch.full((1, 1, 2, 2), 0.5, dtype=torch.float32)
torch.testing.assert_close(match_color_mean(sr, lr), sr, rtol=0.0, atol=0.0)
def test_unclipped_shift_moves_every_channel_mean_to_the_lr_mean(self) -> None:
sr = torch.arange(12, dtype=torch.float32).reshape(1, 1, 3, 4) / 16.0
lr = torch.full((1, 1, 3, 4), 0.5, dtype=torch.float32)
shifted = match_color_mean(sr, lr)
torch.testing.assert_close(shifted.mean(dim=(-2, -1)), lr.mean(dim=(-2, -1)), rtol=0.0, atol=0.0)
difference = shifted - sr
torch.testing.assert_close(
difference, difference.flatten()[0].expand_as(difference), rtol=0.0, atol=0.0
)
def test_computes_in_float32_regardless_of_input_dtype(self) -> None:
sr = torch.tensor([[[[0.25, 0.5], [0.5, 0.75]]]], dtype=torch.bfloat16)
lr = torch.tensor([[[[0.5, 0.5], [0.5, 0.5]]]], dtype=torch.bfloat16)
self.assertEqual(match_color_mean(sr, lr).dtype, torch.float32)
def test_does_not_mutate_its_inputs(self) -> None:
sr = torch.tensor([[[[0.25, 0.75], [0.25, 0.75]]]], dtype=torch.float32)
lr = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]], dtype=torch.float32)
sr_copy, lr_copy = sr.clone(), lr.clone()
match_color_mean(sr, lr)
torch.testing.assert_close(sr, sr_copy, rtol=0.0, atol=0.0)
torch.testing.assert_close(lr, lr_copy, rtol=0.0, atol=0.0)
class ProtocolConstantTests(unittest.TestCase):
def test_inference_protocol_constant_matches_the_verifier(self) -> None:
self.assertEqual(MEAN_ONLY_PROTOCOL, VERIFIER_MEAN_ONLY_PROTOCOL)
def test_documented_formula_is_the_per_channel_mean_shift(self) -> None:
self.assertEqual(COLOR_MATCH_FORMULA, "clamp(sr - mean(sr,H,W) + mean(lr,H,W), 0, 1)")
def test_documented_model_choices(self) -> None:
self.assertEqual(MODEL_CHOICES, ("raesr", "raesr-d"))
class FrontendWeightsTests(unittest.TestCase):
def test_released_frontend_weights_match_the_projector_contract(self) -> None:
from raesr.runtime.frontends.raesr_d import PatchTokenProjector
weights = DELIVERY_ROOT / "weights" / "raesr_d_frontend.pt"
self.assertTrue(weights.is_file(), f"missing frontend weights: {weights}")
state_dict = torch.load(weights, map_location="cpu", weights_only=True)
projector = PatchTokenProjector()
reference = projector.state_dict()
self.assertEqual(sorted(state_dict), sorted(reference))
for key, tensor in reference.items():
self.assertEqual(tuple(state_dict[key].shape), tuple(tensor.shape), key)
projector.load_state_dict(state_dict, strict=True)
def test_frontend_state_dict_contains_no_python_objects(self) -> None:
weights = DELIVERY_ROOT / "weights" / "raesr_d_frontend.pt"
self.assertTrue(weights.is_file(), f"missing frontend weights: {weights}")
state_dict = torch.load(weights, map_location="cpu", weights_only=True)
self.assertTrue(all(torch.is_tensor(value) for value in state_dict.values()))
if __name__ == "__main__":
unittest.main()
|