File size: 2,185 Bytes
6826a45 | 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 | from __future__ import annotations
from pathlib import Path
import torch
from gtcrn_coreai.loader import (
ATTENTION_STATE_SHAPE,
CONVOLUTION_STATE_SHAPE,
INTER_FRAME_STATE_SHAPE,
create_streaming_model,
initial_state,
)
from gtcrn_coreai.model import GTCRN
ROOT = Path(__file__).resolve().parents[1]
def test_initial_state_shapes_and_dtype() -> None:
state = initial_state(dtype=torch.float16)
assert tuple(state[0].shape) == CONVOLUTION_STATE_SHAPE
assert tuple(state[1].shape) == ATTENTION_STATE_SHAPE
assert tuple(state[2].shape) == INTER_FRAME_STATE_SHAPE
assert all(value.dtype == torch.float16 for value in state)
assert all(torch.count_nonzero(value) == 0 for value in state)
def test_functional_wrapper_does_not_mutate_inputs() -> None:
model = create_streaming_model(ROOT / "weights/model_trained_on_dns3.tar")
spectrogram = torch.randn(1, 257, 1, 2)
state = initial_state()
before = tuple(value.clone() for value in state)
with torch.no_grad():
outputs = model(spectrogram, *state)
assert tuple(outputs[0].shape) == (1, 257, 1, 2)
assert all(
torch.equal(actual, expected)
for actual, expected in zip(state, before, strict=True)
)
assert any(torch.count_nonzero(value) > 0 for value in outputs[1:])
def test_streaming_frames_match_offline_model() -> None:
checkpoint_path = ROOT / "weights/model_trained_on_dns3.tar"
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
offline_model = GTCRN().eval()
offline_model.load_state_dict(checkpoint["model"], strict=True)
streaming_model = create_streaming_model(checkpoint_path)
spectrogram = torch.randn(1, 257, 12, 2)
with torch.no_grad():
expected = offline_model(spectrogram)
state = initial_state()
frames = []
for index in range(spectrogram.shape[2]):
outputs = streaming_model(spectrogram[:, :, index : index + 1], *state)
frames.append(outputs[0])
state = outputs[1:]
actual = torch.cat(frames, dim=2)
torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5)
|