Add frozen EgoSpeak Omni residual fusion stage
This commit is contained in:
Xin Wang
169
test_egospeak_omni_residual.py
Normal file
169
test_egospeak_omni_residual.py
Normal file
@@ -0,0 +1,169 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from egospeak_omni_residual import (
|
||||
EgoSpeakOmniResidual,
|
||||
EgoSpeakOmniWindowDataset,
|
||||
MiniROADABase,
|
||||
align_source_predictions,
|
||||
load_easycom_session_arrays,
|
||||
load_frozen_miniroada,
|
||||
miniroad_normalized_loss,
|
||||
)
|
||||
|
||||
|
||||
class EgoSpeakOmniResidualTest(unittest.TestCase):
|
||||
def test_raw_miniroada_state_dict_loads_strictly_and_freezes(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
checkpoint = Path(directory) / "base.pth"
|
||||
source = MiniROADABase(
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
embedding_dim=8,
|
||||
hidden_dim=5,
|
||||
num_horizons=10,
|
||||
)
|
||||
torch.save(source.state_dict(), checkpoint)
|
||||
loaded = load_frozen_miniroada(
|
||||
checkpoint,
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
embedding_dim=8,
|
||||
hidden_dim=5,
|
||||
num_horizons=10,
|
||||
)
|
||||
self.assertFalse(loaded.training)
|
||||
self.assertTrue(
|
||||
all(not parameter.requires_grad for parameter in loaded.parameters())
|
||||
)
|
||||
for key, value in source.state_dict().items():
|
||||
torch.testing.assert_close(loaded.state_dict()[key], value)
|
||||
|
||||
def test_zero_residual_exactly_recovers_frozen_base(self) -> None:
|
||||
torch.manual_seed(7)
|
||||
base = MiniROADABase(
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
embedding_dim=8,
|
||||
hidden_dim=5,
|
||||
num_horizons=10,
|
||||
dropout=0.4,
|
||||
)
|
||||
model = EgoSpeakOmniResidual(
|
||||
base,
|
||||
omni_dim=7,
|
||||
fusion_dim=4,
|
||||
tail_embedding_dim=3,
|
||||
dropout=0.3,
|
||||
)
|
||||
model.train()
|
||||
self.assertFalse(model.base.training)
|
||||
self.assertTrue(
|
||||
all(not parameter.requires_grad for parameter in model.base.parameters())
|
||||
)
|
||||
|
||||
rgb = torch.randn(2, 9, 4)
|
||||
audio = torch.randn(2, 9, 6)
|
||||
omni = torch.randn(2, 7)
|
||||
tail = torch.tensor([0, 9], dtype=torch.long)
|
||||
fused, frozen = model.forward_last(rgb, audio, omni, tail)
|
||||
torch.testing.assert_close(fused, frozen, rtol=0.0, atol=0.0)
|
||||
|
||||
targets = torch.zeros_like(fused)
|
||||
targets[..., 1] = 1
|
||||
optimizer = torch.optim.SGD(model.trainable_parameters(), lr=0.1)
|
||||
loss = miniroad_normalized_loss(fused, targets)
|
||||
loss.backward()
|
||||
self.assertGreater(float(model.residual_output.weight.grad.abs().sum()), 0.0)
|
||||
optimizer.step()
|
||||
updated, frozen_again = model.forward_last(rgb, audio, omni, tail)
|
||||
self.assertGreater(
|
||||
float((updated - frozen_again).abs().max().detach()),
|
||||
0.0,
|
||||
)
|
||||
|
||||
def test_feature_contract_and_window_alignment(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
ego = root / "ego"
|
||||
omni = root / "omni"
|
||||
for subdir in (
|
||||
"rgb_kinetics_resnet50",
|
||||
"audio_wav2vec2",
|
||||
"target_perframe",
|
||||
):
|
||||
(ego / subdir).mkdir(parents=True, exist_ok=True)
|
||||
omni.mkdir()
|
||||
steps = 15
|
||||
rgb = np.arange(steps * 4, dtype=np.float32).reshape(steps, 4)
|
||||
audio = np.arange(steps * 6, dtype=np.float32).reshape(steps, 6)
|
||||
target = np.zeros((steps, 3), dtype=np.float32)
|
||||
target[:, 0] = 1
|
||||
np.save(ego / "rgb_kinetics_resnet50" / "Session_4.npy", rgb)
|
||||
np.save(ego / "audio_wav2vec2" / "Session_4.npy", audio)
|
||||
np.save(ego / "target_perframe" / "Session_4.npy", target)
|
||||
torch.save(
|
||||
{
|
||||
"mode": "open-stream-replay",
|
||||
"states": torch.arange(steps * 7, dtype=torch.float32).view(
|
||||
steps, 7
|
||||
),
|
||||
"times": (torch.arange(steps, dtype=torch.float64) + 1) / 5,
|
||||
"label_indices": torch.arange(steps, dtype=torch.long),
|
||||
"start": 0.0,
|
||||
"config": {
|
||||
"step_seconds": 0.2,
|
||||
"native_block_seconds": 2.0,
|
||||
},
|
||||
},
|
||||
omni / "Session_4.pt",
|
||||
)
|
||||
session = load_easycom_session_arrays(
|
||||
ego,
|
||||
omni,
|
||||
"Session_4",
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
)
|
||||
self.assertEqual(
|
||||
session.tail_bins.tolist(), list(range(10)) + list(range(5))
|
||||
)
|
||||
dataset = EgoSpeakOmniWindowDataset(
|
||||
[session],
|
||||
window_size=4,
|
||||
stride=1,
|
||||
num_horizons=3,
|
||||
seed=0,
|
||||
)
|
||||
first_index = dataset.indices.index((0, 0))
|
||||
rgb_window, audio_window, state, tail, future, _, source = dataset[
|
||||
first_index
|
||||
]
|
||||
self.assertEqual(int(source), 0)
|
||||
self.assertEqual(int(tail), 0)
|
||||
torch.testing.assert_close(rgb_window[-1], torch.from_numpy(rgb[0]))
|
||||
self.assertEqual(int(torch.count_nonzero(rgb_window[:-1])), 0)
|
||||
torch.testing.assert_close(audio_window[-1], torch.from_numpy(audio[0]))
|
||||
torch.testing.assert_close(state, session.omni[0])
|
||||
torch.testing.assert_close(future, torch.from_numpy(target[1:4]))
|
||||
|
||||
def test_source_prediction_alignment(self) -> None:
|
||||
target_steps = 14
|
||||
raw = np.arange(
|
||||
(target_steps - 10) * 10 * 3,
|
||||
dtype=np.float32,
|
||||
).reshape(target_steps - 10, 10, 3)
|
||||
pred, valid = align_source_predictions(raw, target_steps)
|
||||
np.testing.assert_array_equal(pred[1, :, 0], raw[0, 0, :])
|
||||
np.testing.assert_array_equal(pred[13, :, 9], raw[3, 9, :])
|
||||
self.assertTrue(valid[1, 0])
|
||||
self.assertFalse(valid[0, 0])
|
||||
self.assertTrue(valid[13, 9])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user