wehub-resource-sync
624 lines
20 KiB
Python
624 lines
20 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unit tests for TurboQuant serde skeleton.
|
|
"""
|
|
|
|
# Standard
|
|
from pathlib import Path
|
|
from typing import cast
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.native_storage_ops import Bitmap
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.config import (
|
|
EvictionConfig,
|
|
L1ManagerConfig,
|
|
L1MemoryManagerConfig,
|
|
StorageManagerConfig,
|
|
)
|
|
from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig
|
|
from lmcache.v1.distributed.l2_adapters.fs_l2_adapter import FSL2AdapterConfig
|
|
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig
|
|
from lmcache.v1.distributed.serde import (
|
|
SerdeConfig,
|
|
create_serde_processor,
|
|
get_registered_serde_types,
|
|
)
|
|
from lmcache.v1.distributed.serde.turboquant import (
|
|
TurboQuantSerdeConfig,
|
|
TurboQuantSerializer,
|
|
)
|
|
from lmcache.v1.distributed.storage_manager import StorageManager
|
|
from lmcache.v1.memory_management import MemoryObj
|
|
|
|
|
|
def test_turboquant_registered() -> None:
|
|
assert "turboquant" in get_registered_serde_types()
|
|
|
|
|
|
def test_create_turboquant_serde_processor() -> None:
|
|
processor = create_serde_processor(
|
|
SerdeConfig(
|
|
type="turboquant",
|
|
kwargs={
|
|
"preset": "turboquant_k8v4",
|
|
"head_dim": 128,
|
|
"block_size": 16,
|
|
"max_workers": 1,
|
|
},
|
|
)
|
|
)
|
|
processor.close()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
(
|
|
"preset",
|
|
"key_fp8",
|
|
"key_quant_bits",
|
|
"key_mse_bits",
|
|
"value_quant_bits",
|
|
"norm_correction",
|
|
"key_packed_size",
|
|
"value_packed_size",
|
|
"slot_size",
|
|
"slot_size_aligned",
|
|
),
|
|
[
|
|
("turboquant_k8v4", True, 8, 0, 4, False, 128, 68, 196, 196),
|
|
("turboquant_4bit_nc", False, 4, 4, 4, True, 66, 68, 134, 134),
|
|
("turboquant_k3v4_nc", False, 3, 3, 4, True, 50, 68, 118, 118),
|
|
("turboquant_3bit_nc", False, 3, 3, 3, True, 50, 52, 102, 102),
|
|
],
|
|
)
|
|
def test_turboquant_config_sizes_head_dim_128(
|
|
preset: str,
|
|
key_fp8: bool,
|
|
key_quant_bits: int,
|
|
key_mse_bits: int,
|
|
value_quant_bits: int,
|
|
norm_correction: bool,
|
|
key_packed_size: int,
|
|
value_packed_size: int,
|
|
slot_size: int,
|
|
slot_size_aligned: int,
|
|
) -> None:
|
|
cfg = TurboQuantSerdeConfig(
|
|
preset=preset,
|
|
head_dim=128,
|
|
block_size=16,
|
|
)
|
|
|
|
assert cfg.key_fp8 is key_fp8
|
|
assert cfg.key_quant_bits == key_quant_bits
|
|
assert cfg.key_mse_bits == key_mse_bits
|
|
assert cfg.value_quant_bits == value_quant_bits
|
|
assert cfg.norm_correction is norm_correction
|
|
assert cfg.key_packed_size == key_packed_size
|
|
assert cfg.value_packed_size == value_packed_size
|
|
assert cfg.slot_size == slot_size
|
|
assert cfg.slot_size_aligned == slot_size_aligned
|
|
|
|
|
|
def test_turboquant_config_rejects_invalid_preset() -> None:
|
|
cfg = TurboQuantSerdeConfig(
|
|
preset="turboquant_invalid",
|
|
head_dim=128,
|
|
block_size=16,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Unsupported TurboQuant preset"):
|
|
_ = cfg.key_quant_bits
|
|
|
|
|
|
def test_estimate_serialized_size_k8v4() -> None:
|
|
cfg = TurboQuantSerdeConfig(
|
|
preset="turboquant_k8v4",
|
|
head_dim=128,
|
|
block_size=16,
|
|
)
|
|
serializer = TurboQuantSerializer(cfg)
|
|
|
|
# LMCache KV layout: [2, num_layers, num_tokens, hidden_dim]
|
|
# hidden_dim = num_heads * head_dim = 4 * 128 = 512
|
|
layout = MemoryLayoutDesc(
|
|
shapes=[torch.Size([2, 3, 20, 512])],
|
|
dtypes=[torch.bfloat16],
|
|
)
|
|
|
|
# num_layers = 3
|
|
# default skip_first_layers=2 and skip_last_layers=2 leaves no middle
|
|
# layers to compress, so all layers are stored as raw bfloat16 KV.
|
|
expected = 2 * 3 * 20 * 512 * torch.bfloat16.itemsize
|
|
|
|
assert serializer.estimate_serialized_size(layout) == expected
|
|
|
|
|
|
def test_estimate_serialized_size_rejects_invalid_kv_size() -> None:
|
|
cfg = TurboQuantSerdeConfig(
|
|
preset="turboquant_k8v4",
|
|
head_dim=128,
|
|
block_size=16,
|
|
)
|
|
serializer = TurboQuantSerializer(cfg)
|
|
|
|
layout = MemoryLayoutDesc(
|
|
shapes=[torch.Size([1, 3, 20, 512])],
|
|
dtypes=[torch.bfloat16],
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="kv_size=2"):
|
|
serializer.estimate_serialized_size(layout)
|
|
|
|
|
|
def test_estimate_serialized_size_rejects_bad_head_dim() -> None:
|
|
cfg = TurboQuantSerdeConfig(
|
|
preset="turboquant_k8v4",
|
|
head_dim=128,
|
|
block_size=16,
|
|
)
|
|
serializer = TurboQuantSerializer(cfg)
|
|
|
|
layout = MemoryLayoutDesc(
|
|
shapes=[torch.Size([2, 3, 20, 500])],
|
|
dtypes=[torch.bfloat16],
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="must be divisible"):
|
|
serializer.estimate_serialized_size(layout)
|
|
|
|
|
|
# =============================================================================
|
|
# GPU E2E test: StorageManager + SerdeL2AdapterWrapper + MockL2Adapter
|
|
# =============================================================================
|
|
|
|
|
|
def _make_turboquant_object_key(chunk_id: int) -> ObjectKey:
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
|
|
model_name="turboquant_test_model",
|
|
kv_rank=0,
|
|
)
|
|
|
|
|
|
def _make_turboquant_layout() -> MemoryLayoutDesc:
|
|
# TurboQuant serde currently expects [2, num_layers, num_tokens, hidden_dim].
|
|
# hidden_dim = num_heads * head_dim = 4 * 128 = 512.
|
|
return MemoryLayoutDesc(
|
|
shapes=[torch.Size([2, 2, 32, 512])],
|
|
dtypes=[torch.bfloat16],
|
|
)
|
|
|
|
|
|
def _wait_for_condition(
|
|
predicate,
|
|
timeout: float = 20.0,
|
|
poll_interval: float = 0.05,
|
|
) -> bool:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(poll_interval)
|
|
return False
|
|
|
|
|
|
def _wait_for_prefetch_status(
|
|
sm: StorageManager,
|
|
handle,
|
|
timeout: float = 20.0,
|
|
poll_interval: float = 0.05,
|
|
) -> Bitmap | None:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
result = sm.query_prefetch_status(handle)
|
|
if result is not None:
|
|
return result
|
|
time.sleep(poll_interval)
|
|
return None
|
|
|
|
|
|
def _finish_read_prefetched_until_clean(
|
|
sm: StorageManager,
|
|
keys: list[ObjectKey],
|
|
timeout: float = 30.0,
|
|
) -> None:
|
|
"""Release prefetched temporary objects and wait for L1 cleanup.
|
|
|
|
StorageManager / serde wrapper paths may hold more than one read lock
|
|
on temporary prefetched objects. Release repeatedly until L1 is clean.
|
|
"""
|
|
for _ in range(4):
|
|
sm.finish_read_prefetched(keys)
|
|
ok = _wait_for_condition(
|
|
lambda: (
|
|
sm.report_status()["l1_manager"]["memory_used_bytes"] == 0
|
|
and sm.report_status()["l1_manager"]["total_object_count"] == 0
|
|
and sm.report_status()["l1_manager"]["read_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["write_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["temporary_count"] == 0
|
|
),
|
|
timeout=timeout,
|
|
)
|
|
if ok:
|
|
return
|
|
|
|
raise AssertionError(f"L1 memory not released: {sm.report_status()['l1_manager']}")
|
|
|
|
|
|
def _make_turboquant_storage_manager(preset: str) -> StorageManager:
|
|
adapter_cfg = MockL2AdapterConfig(
|
|
max_size_gb=0.1,
|
|
mock_bandwidth_gb=10.0,
|
|
)
|
|
adapter_cfg.serde_config = SerdeConfig(
|
|
type="turboquant",
|
|
kwargs={
|
|
"preset": preset,
|
|
"head_dim": 128,
|
|
"block_size": 16,
|
|
"max_workers": 1,
|
|
},
|
|
)
|
|
|
|
cfg = StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=256 * 1024 * 1024,
|
|
use_lazy=torch.cuda.is_available(),
|
|
init_size_in_bytes=64 * 1024 * 1024,
|
|
),
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="LRU"),
|
|
l2_adapter_config=L2AdaptersConfig(adapters=[adapter_cfg]),
|
|
)
|
|
return StorageManager(cfg)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
|
|
@pytest.mark.parametrize(
|
|
("preset", "corr_lower_bound"),
|
|
[
|
|
("turboquant_k8v4", 0.95),
|
|
("turboquant_4bit_nc", 0.90),
|
|
("turboquant_k3v4_nc", 0.85),
|
|
("turboquant_3bit_nc", 0.80),
|
|
],
|
|
)
|
|
def test_turboquant_storage_manager_roundtrip(
|
|
preset: str,
|
|
corr_lower_bound: float,
|
|
) -> None:
|
|
"""Store/load through StorageManager with TurboQuant serde.
|
|
|
|
This verifies:
|
|
reserve_write -> finish_write -> StoreController
|
|
-> SerdeL2AdapterWrapper serialize -> MockL2Adapter store
|
|
-> clear L1 -> prefetch -> MockL2Adapter load
|
|
-> SerdeL2AdapterWrapper deserialize -> read_prefetched_results.
|
|
"""
|
|
sm = _make_turboquant_storage_manager(preset)
|
|
layout = _make_turboquant_layout()
|
|
keys = [_make_turboquant_object_key(i) for i in range(3)]
|
|
|
|
try:
|
|
ret = sm.reserve_write(keys, layout, mode="new")
|
|
assert len(ret) == len(keys), f"reserve_write got {len(ret)} / {len(keys)}"
|
|
|
|
original_by_key = {}
|
|
for i, key in enumerate(keys):
|
|
obj = ret[key]
|
|
assert obj.tensor is not None
|
|
|
|
torch.manual_seed(1234 + i)
|
|
data = torch.randn(
|
|
obj.tensor.shape,
|
|
dtype=obj.tensor.dtype,
|
|
device=obj.tensor.device,
|
|
)
|
|
data = data + float(i)
|
|
|
|
obj.tensor.copy_(data)
|
|
original_by_key[key] = data.detach().clone()
|
|
|
|
sm.finish_write(list(ret.keys()))
|
|
|
|
# Wait until both the store controller and the serde wrapper cleanup
|
|
# have released temporary objects and locks.
|
|
ok = _wait_for_condition(
|
|
lambda: (
|
|
sm.report_status()["store_controller"]["in_flight_task_count"] == 0
|
|
and sm.report_status()["store_controller"]["pending_keys_count"] == 0
|
|
and sm.report_status()["l1_manager"]["write_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["read_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["temporary_count"] == 0
|
|
),
|
|
timeout=120.0,
|
|
)
|
|
assert ok, "Store to L2 did not fully complete"
|
|
|
|
sm.clear()
|
|
ok = _wait_for_condition(
|
|
lambda: (
|
|
sm.report_status()["l1_manager"]["total_object_count"] == 0
|
|
and sm.report_status()["l1_manager"]["memory_used_bytes"] == 0
|
|
and sm.report_status()["l1_manager"]["read_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["write_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["temporary_count"] == 0
|
|
),
|
|
timeout=120.0,
|
|
)
|
|
assert ok, f"L1 not cleared: {sm.report_status()['l1_manager']}"
|
|
|
|
handle = sm.submit_prefetch_task(keys, layout)
|
|
hit_bitmap = _wait_for_prefetch_status(sm, handle, timeout=120.0)
|
|
assert hit_bitmap is not None
|
|
hits = hit_bitmap.count_leading_ones()
|
|
assert hits == len(keys), f"Expected {len(keys)} hits, got {hits}"
|
|
|
|
with sm.read_prefetched_results(keys) as objs:
|
|
assert objs is not None
|
|
assert len(objs) == len(keys)
|
|
|
|
for key, obj in zip(keys, objs, strict=True):
|
|
assert obj.tensor is not None
|
|
|
|
recovered = obj.tensor
|
|
original = original_by_key[key]
|
|
|
|
orig_f = original.float().flatten()
|
|
rec_f = recovered.float().flatten()
|
|
|
|
corr = torch.corrcoef(torch.stack([orig_f, rec_f]))[0, 1].item()
|
|
mae = torch.mean(torch.abs(orig_f - rec_f)).item()
|
|
mse = torch.mean((orig_f - rec_f) ** 2).item()
|
|
|
|
assert corr > corr_lower_bound, (
|
|
f"low corr for preset={preset}, key={key}: "
|
|
f"corr={corr}, mae={mae}, mse={mse}"
|
|
)
|
|
|
|
_finish_read_prefetched_until_clean(sm, keys)
|
|
finally:
|
|
sm.close()
|
|
|
|
|
|
# =============================================================================
|
|
# GPU direct test: TurboQuantSerializer + TurboQuantDeserializer
|
|
# =============================================================================
|
|
|
|
|
|
class _FakeMemoryObj:
|
|
def __init__(self, tensor: torch.Tensor):
|
|
self.tensor = tensor
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
|
|
@pytest.mark.parametrize(
|
|
("preset", "expected_ratio_lower_bound", "corr_lower_bound"),
|
|
[
|
|
("turboquant_k8v4", 2.60, 0.95),
|
|
("turboquant_4bit_nc", 3.75, 0.90),
|
|
("turboquant_k3v4_nc", 4.20, 0.85),
|
|
("turboquant_3bit_nc", 4.85, 0.80),
|
|
],
|
|
)
|
|
def test_turboquant_direct_roundtrip_cuda(
|
|
preset: str,
|
|
expected_ratio_lower_bound: float,
|
|
corr_lower_bound: float,
|
|
) -> None:
|
|
"""Direct GPU round-trip through TurboQuant serializer/deserializer."""
|
|
# First Party
|
|
from lmcache.v1.distributed.serde.turboquant import TurboQuantDeserializer
|
|
|
|
device = torch.device("cuda:0")
|
|
dtype = torch.float16
|
|
|
|
# LMCache KV layout: [2, num_layers, num_tokens, hidden_dim]
|
|
num_layers = 4
|
|
num_tokens = 128
|
|
num_heads = 8
|
|
head_dim = 128
|
|
hidden_dim = num_heads * head_dim
|
|
|
|
cfg = TurboQuantSerdeConfig(
|
|
preset=preset,
|
|
head_dim=head_dim,
|
|
block_size=16,
|
|
skip_first_layers=0,
|
|
skip_last_layers=0,
|
|
)
|
|
|
|
torch.manual_seed(2026)
|
|
shape = torch.Size([2, num_layers, num_tokens, hidden_dim])
|
|
original = torch.randn(shape, dtype=dtype, device=device)
|
|
|
|
serializer = TurboQuantSerializer(cfg)
|
|
deserializer = TurboQuantDeserializer(cfg)
|
|
|
|
layout = MemoryLayoutDesc(shapes=[shape], dtypes=[dtype])
|
|
n_bytes = serializer.estimate_serialized_size(layout)
|
|
|
|
compressed = torch.empty(n_bytes, dtype=torch.uint8, device=device)
|
|
recovered = torch.empty_like(original)
|
|
|
|
written = serializer.serialize(
|
|
cast(MemoryObj, _FakeMemoryObj(original)),
|
|
cast(MemoryObj, _FakeMemoryObj(compressed)),
|
|
)
|
|
assert written == n_bytes
|
|
|
|
deserializer.deserialize(
|
|
cast(MemoryObj, _FakeMemoryObj(compressed)),
|
|
cast(MemoryObj, _FakeMemoryObj(recovered)),
|
|
)
|
|
|
|
orig_f = original.float().flatten()
|
|
rec_f = recovered.float().flatten()
|
|
|
|
corr = torch.corrcoef(torch.stack([orig_f, rec_f]))[0, 1].item()
|
|
mae = torch.mean(torch.abs(orig_f - rec_f)).item()
|
|
mse = torch.mean((orig_f - rec_f) ** 2).item()
|
|
|
|
original_bytes = original.numel() * original.element_size()
|
|
ratio = original_bytes / n_bytes
|
|
|
|
assert ratio >= expected_ratio_lower_bound
|
|
assert corr > corr_lower_bound, (
|
|
f"low corr for preset={preset}: corr={corr}, mae={mae}, mse={mse}"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# GPU E2E test: StorageManager + SerdeL2AdapterWrapper + FSL2Adapter
|
|
# =============================================================================
|
|
|
|
|
|
def _make_turboquant_fs_storage_manager(
|
|
base_path: str,
|
|
preset: str,
|
|
) -> StorageManager:
|
|
adapter_cfg = FSL2AdapterConfig(
|
|
base_path=base_path,
|
|
relative_tmp_dir="tmp",
|
|
use_odirect=False,
|
|
)
|
|
adapter_cfg.serde_config = SerdeConfig(
|
|
type="turboquant",
|
|
kwargs={
|
|
"preset": preset,
|
|
"head_dim": 128,
|
|
"block_size": 16,
|
|
"max_workers": 1,
|
|
},
|
|
)
|
|
|
|
cfg = StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=256 * 1024 * 1024,
|
|
use_lazy=torch.cuda.is_available(),
|
|
init_size_in_bytes=64 * 1024 * 1024,
|
|
),
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="LRU"),
|
|
l2_adapter_config=L2AdaptersConfig(adapters=[adapter_cfg]),
|
|
)
|
|
return StorageManager(cfg)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
|
|
@pytest.mark.parametrize(
|
|
("preset", "corr_lower_bound"),
|
|
[
|
|
("turboquant_k8v4", 0.95),
|
|
("turboquant_4bit_nc", 0.90),
|
|
("turboquant_k3v4_nc", 0.85),
|
|
("turboquant_3bit_nc", 0.80),
|
|
],
|
|
)
|
|
def test_turboquant_fs_storage_manager_roundtrip(
|
|
preset: str,
|
|
corr_lower_bound: float,
|
|
) -> None:
|
|
"""Store/load through FSL2Adapter with TurboQuant serde."""
|
|
base_dir = tempfile.mkdtemp(prefix="lmcache_turboquant_fs_")
|
|
sm = None
|
|
|
|
try:
|
|
sm = _make_turboquant_fs_storage_manager(base_dir, preset)
|
|
layout = _make_turboquant_layout()
|
|
keys = [_make_turboquant_object_key(i) for i in range(3)]
|
|
|
|
ret = sm.reserve_write(keys, layout, mode="new")
|
|
assert len(ret) == len(keys), f"reserve_write got {len(ret)} / {len(keys)}"
|
|
|
|
original_by_key = {}
|
|
for i, key in enumerate(keys):
|
|
obj = ret[key]
|
|
assert obj.tensor is not None
|
|
|
|
torch.manual_seed(5678 + i)
|
|
data = torch.randn(
|
|
obj.tensor.shape,
|
|
dtype=obj.tensor.dtype,
|
|
device=obj.tensor.device,
|
|
)
|
|
data = data + float(i)
|
|
|
|
obj.tensor.copy_(data)
|
|
original_by_key[key] = data.detach().clone()
|
|
|
|
sm.finish_write(list(ret.keys()))
|
|
|
|
ok = _wait_for_condition(
|
|
lambda: (
|
|
sm.report_status()["store_controller"]["in_flight_task_count"] == 0
|
|
and sm.report_status()["store_controller"]["pending_keys_count"] == 0
|
|
and sm.report_status()["l1_manager"]["write_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["read_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["temporary_count"] == 0
|
|
),
|
|
timeout=120.0,
|
|
)
|
|
assert ok, "Store to FS L2 did not fully complete"
|
|
|
|
stored_files = [p for p in Path(base_dir).rglob("*") if p.is_file()]
|
|
assert len(stored_files) >= len(keys)
|
|
|
|
sm.clear()
|
|
ok = _wait_for_condition(
|
|
lambda: (
|
|
sm.report_status()["l1_manager"]["total_object_count"] == 0
|
|
and sm.report_status()["l1_manager"]["memory_used_bytes"] == 0
|
|
and sm.report_status()["l1_manager"]["read_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["write_locked_count"] == 0
|
|
and sm.report_status()["l1_manager"]["temporary_count"] == 0
|
|
),
|
|
timeout=120.0,
|
|
)
|
|
assert ok, f"L1 not cleared: {sm.report_status()['l1_manager']}"
|
|
|
|
handle = sm.submit_prefetch_task(keys, layout)
|
|
hit_bitmap = _wait_for_prefetch_status(sm, handle, timeout=120.0)
|
|
assert hit_bitmap is not None
|
|
hits = hit_bitmap.count_leading_ones()
|
|
assert hits == len(keys), f"Expected {len(keys)} hits, got {hits}"
|
|
|
|
with sm.read_prefetched_results(keys) as objs:
|
|
assert objs is not None
|
|
assert len(objs) == len(keys)
|
|
|
|
for key, obj in zip(keys, objs, strict=True):
|
|
assert obj.tensor is not None
|
|
|
|
recovered = obj.tensor
|
|
original = original_by_key[key]
|
|
|
|
orig_f = original.float().flatten()
|
|
rec_f = recovered.float().flatten()
|
|
|
|
corr = torch.corrcoef(torch.stack([orig_f, rec_f]))[0, 1].item()
|
|
mae = torch.mean(torch.abs(orig_f - rec_f)).item()
|
|
mse = torch.mean((orig_f - rec_f) ** 2).item()
|
|
|
|
assert corr > corr_lower_bound, (
|
|
f"low corr for preset={preset}, key={key}: "
|
|
f"corr={corr}, mae={mae}, mse={mse}"
|
|
)
|
|
|
|
_finish_read_prefetched_until_clean(sm, keys)
|
|
|
|
finally:
|
|
if sm is not None:
|
|
sm.close()
|
|
shutil.rmtree(base_dir, ignore_errors=True)
|