| import unittest |
|
|
| import torch |
|
|
| from losses.video_nystrom_mmd import ( |
| VideoNystromMMD, |
| build_video_nystrom_reference, |
| rbf_frobenius_kernel, |
| ) |
|
|
|
|
| class TestVideoNystromPartialMMD(unittest.TestCase): |
| def _pack(self): |
| reference = torch.randn( |
| 12, |
| 1, |
| 2, |
| generator=torch.Generator().manual_seed(9100), |
| ) |
| return build_video_nystrom_reference( |
| reference, |
| num_landmarks=5, |
| metadata={ |
| "feature": { |
| "name": "tiny-partial-video", |
| "matrix_shape": [1, 2], |
| }, |
| "reference": {"dataset": "partial-unit-test"}, |
| }, |
| bandwidth=1.6, |
| kmeans_iterations=4, |
| ridge=1.0e-5, |
| chunk_size=3, |
| ) |
|
|
| def _matched_objective(self, *, chunk_size: int = 37): |
| return VideoNystromMMD( |
| self._pack(), |
| kernel_chunk_size=chunk_size, |
| attraction_weight=2.0, |
| current_teacher_attraction_enabled=True, |
| current_teacher_attraction_weight=1.0, |
| ) |
|
|
| def test_128_by_256_formula_denominators_and_suffix_only_gradient(self): |
| objective = self._matched_objective() |
| generator = torch.Generator().manual_seed(9101) |
| all_students = torch.randn( |
| 256, 1, 2, generator=generator, requires_grad=True) |
| gradient_students = all_students[-128:] |
| teacher = torch.randn( |
| 128, 1, 2, generator=generator, requires_grad=True) |
|
|
| loss, metadata = objective( |
| gradient_students=gradient_students, |
| student_reference=all_students, |
| current_teacher=teacher, |
| stage=False, |
| ) |
|
|
| kernel_gs = rbf_frobenius_kernel( |
| gradient_students, |
| all_students.detach(), |
| objective.sigma, |
| ) |
| fixed_attraction = ( |
| rbf_frobenius_kernel( |
| gradient_students, |
| objective.landmarks, |
| objective.sigma, |
| ) |
| @ objective.alpha |
| ).mean() |
| kernel_gt = rbf_frobenius_kernel( |
| gradient_students, |
| teacher.detach(), |
| objective.sigma, |
| ) |
| kernel_tt = rbf_frobenius_kernel( |
| teacher.detach(), |
| teacher.detach(), |
| objective.sigma, |
| ) |
| expected = ( |
| kernel_gs.sum() / float(128 * 256) |
| - 2.0 |
| * ( |
| (2.0 / 3.0) * fixed_attraction |
| + (1.0 / 3.0) * kernel_gt.sum() / float(128 * 128) |
| ) |
| + (2.0 / 3.0) * objective.k_rr |
| + (1.0 / 3.0) * kernel_tt.sum() / float(128 * 128) |
| ) |
|
|
| torch.testing.assert_close( |
| loss, expected, rtol=2.0e-5, atol=2.0e-6) |
| torch.testing.assert_close( |
| metadata["repulsion"], |
| (kernel_gs.sum() / float(128 * 256)).detach(), |
| rtol=2.0e-5, |
| atol=2.0e-6, |
| ) |
| self.assertEqual(int(metadata["global_samples"]), 256) |
| self.assertEqual(int(metadata["gradient_student_samples"]), 128) |
| self.assertEqual(int(metadata["student_reference_samples"]), 256) |
| self.assertEqual(int(metadata["current_teacher_samples"]), 128) |
| self.assertEqual(int(metadata["repulsion_denominator"]), 32768) |
| self.assertEqual( |
| int(metadata["student_repulsion_denominator"]), 32768) |
| self.assertEqual( |
| int(metadata["current_teacher_attraction_denominator"]), 16384) |
| self.assertEqual( |
| int(metadata["current_teacher_self_kernel_denominator"]), 16384) |
| self.assertEqual( |
| int(metadata["fixed_teacher_attraction_denominator"]), 128) |
| self.assertEqual( |
| int(metadata["partial_student_reference_enabled"]), 1) |
| self.assertEqual(int(metadata["student_reference_detached"]), 1) |
| self.assertNotIn("full_student_partial_gradient", metadata) |
| self.assertAlmostEqual( |
| float(metadata[ |
| "fixed_teacher_attraction_normalized_weight"]), |
| 2.0 / 3.0, |
| places=7, |
| ) |
| self.assertAlmostEqual( |
| float(metadata[ |
| "current_teacher_attraction_normalized_weight"]), |
| 1.0 / 3.0, |
| places=7, |
| ) |
|
|
| actual_student_gradient, actual_teacher_gradient = ( |
| torch.autograd.grad( |
| loss, |
| (all_students, teacher), |
| allow_unused=True, |
| retain_graph=True, |
| ) |
| ) |
| expected_student_gradient, = torch.autograd.grad( |
| expected, all_students) |
| torch.testing.assert_close( |
| actual_student_gradient, |
| expected_student_gradient, |
| rtol=3.0e-5, |
| atol=3.0e-6, |
| ) |
| self.assertIsNone(actual_teacher_gradient) |
| torch.testing.assert_close( |
| actual_student_gradient[:128], |
| torch.zeros_like(actual_student_gradient[:128]), |
| rtol=0.0, |
| atol=0.0, |
| ) |
| self.assertGreater( |
| float(actual_student_gradient[-128:].norm()), 0.0) |
| self.assertTrue(torch.isfinite(actual_student_gradient).all()) |
|
|
| def test_full_student_partial_gradient_formula_and_denominators(self): |
| objective = self._matched_objective() |
| generator = torch.Generator().manual_seed(9105) |
| student_reference = torch.randn( |
| 256, 1, 2, generator=generator, requires_grad=True) |
| gradient_student_indices = torch.cat([ |
| torch.arange(rank * 64 + 32, rank * 64 + 64) |
| for rank in range(4) |
| ]) |
| gradient_students = ( |
| student_reference.detach() |
| .index_select(0, gradient_student_indices) |
| .clone() |
| .requires_grad_(True) |
| ) |
| teacher = torch.randn( |
| 128, 1, 2, generator=generator, requires_grad=True) |
|
|
| loss, metadata = objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| gradient_student_indices=gradient_student_indices, |
| stage=False, |
| ) |
|
|
| full_students = student_reference.detach().index_copy( |
| 0, gradient_student_indices, gradient_students) |
| kernel_ss = rbf_frobenius_kernel( |
| full_students, full_students, objective.sigma) |
| fixed_attraction = ( |
| rbf_frobenius_kernel( |
| full_students, objective.landmarks, objective.sigma) |
| @ objective.alpha |
| ).mean() |
| kernel_st = rbf_frobenius_kernel( |
| full_students, teacher.detach(), objective.sigma) |
| kernel_tt = rbf_frobenius_kernel( |
| teacher.detach(), teacher.detach(), objective.sigma) |
| expected = ( |
| kernel_ss.sum() / float(256 * 256) |
| - 2.0 |
| * ( |
| (2.0 / 3.0) * fixed_attraction |
| + (1.0 / 3.0) * kernel_st.sum() / float(256 * 128) |
| ) |
| + (2.0 / 3.0) * objective.k_rr |
| + (1.0 / 3.0) * kernel_tt.sum() / float(128 * 128) |
| ) |
|
|
| torch.testing.assert_close( |
| loss, expected, rtol=2.0e-5, atol=2.0e-6) |
| self.assertEqual(int(metadata["global_samples"]), 256) |
| self.assertEqual(int(metadata["gradient_student_samples"]), 128) |
| self.assertEqual(int(metadata["student_reference_samples"]), 256) |
| self.assertEqual(int(metadata["current_teacher_samples"]), 128) |
| self.assertEqual(int(metadata["repulsion_denominator"]), 65536) |
| self.assertEqual( |
| int(metadata["student_repulsion_denominator"]), 65536) |
| self.assertEqual( |
| int(metadata["fixed_teacher_attraction_denominator"]), 256) |
| self.assertEqual( |
| int(metadata["current_teacher_attraction_denominator"]), |
| 32768, |
| ) |
| self.assertEqual( |
| int(metadata["current_teacher_self_kernel_denominator"]), |
| 16384, |
| ) |
| self.assertEqual( |
| int(metadata["full_student_partial_gradient"]), 1) |
| self.assertEqual( |
| int(metadata["student_repulsion_gradient_sides"]), 2) |
|
|
| actual_gradient, reference_gradient, teacher_gradient = ( |
| torch.autograd.grad( |
| loss, |
| (gradient_students, student_reference, teacher), |
| allow_unused=True, |
| retain_graph=True, |
| ) |
| ) |
| expected_gradient, = torch.autograd.grad( |
| expected, gradient_students, retain_graph=True) |
| torch.testing.assert_close( |
| actual_gradient, |
| expected_gradient, |
| rtol=3.0e-5, |
| atol=3.0e-6, |
| ) |
| self.assertIsNone(reference_gradient) |
| self.assertIsNone(teacher_gradient) |
| self.assertGreater(float(actual_gradient.norm()), 0.0) |
| self.assertTrue(torch.isfinite(actual_gradient).all()) |
|
|
| full_repulsion_gradient, = torch.autograd.grad( |
| kernel_ss.sum() / float(256 * 256), |
| gradient_students, |
| retain_graph=True, |
| ) |
| rectangular_repulsion = rbf_frobenius_kernel( |
| gradient_students, student_reference.detach(), objective.sigma |
| ).sum() / float(128 * 256) |
| rectangular_repulsion_gradient, = torch.autograd.grad( |
| rectangular_repulsion, gradient_students) |
| torch.testing.assert_close( |
| full_repulsion_gradient, |
| rectangular_repulsion_gradient, |
| rtol=3.0e-5, |
| atol=3.0e-6, |
| ) |
|
|
| def test_inputs_are_detached_and_converted_to_fp32_for_kernels(self): |
| objective = self._matched_objective(chunk_size=2) |
| generator = torch.Generator().manual_seed(9102) |
| gradient_students = torch.randn( |
| 3, 1, 2, generator=generator, requires_grad=True) |
| student_reference = torch.randn( |
| 5, |
| 1, |
| 2, |
| generator=generator, |
| dtype=torch.float64, |
| requires_grad=True, |
| ) |
| teacher = torch.randn( |
| 3, |
| 1, |
| 2, |
| generator=generator, |
| dtype=torch.float64, |
| requires_grad=True, |
| ) |
|
|
| loss, _ = objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| stage=False, |
| ) |
| student_gradient, reference_gradient, teacher_gradient = ( |
| torch.autograd.grad( |
| loss, |
| (gradient_students, student_reference, teacher), |
| allow_unused=True, |
| ) |
| ) |
| self.assertEqual(loss.dtype, torch.float32) |
| self.assertIsNotNone(student_gradient) |
| self.assertIsNone(reference_gradient) |
| self.assertIsNone(teacher_gradient) |
|
|
| def test_legacy_interface_preserves_full_symmetric_formula_and_schema(self): |
| objective = self._matched_objective(chunk_size=2) |
| generator = torch.Generator().manual_seed(9103) |
| generated = torch.randn( |
| 4, 1, 2, generator=generator, requires_grad=True) |
| teacher = torch.randn(4, 1, 2, generator=generator) |
|
|
| loss, metadata = objective( |
| generated, |
| current_teacher=teacher, |
| stage=False, |
| ) |
| kernel_gg = rbf_frobenius_kernel( |
| generated, generated, objective.sigma) |
| fixed_attraction = ( |
| rbf_frobenius_kernel( |
| generated, objective.landmarks, objective.sigma) |
| @ objective.alpha |
| ).mean() |
| kernel_gt = rbf_frobenius_kernel( |
| generated, teacher, objective.sigma) |
| kernel_tt = rbf_frobenius_kernel( |
| teacher, teacher, objective.sigma) |
| expected = ( |
| kernel_gg.mean() |
| - 2.0 |
| * ( |
| (2.0 / 3.0) * fixed_attraction |
| + (1.0 / 3.0) * kernel_gt.mean() |
| ) |
| + (2.0 / 3.0) * objective.k_rr |
| + (1.0 / 3.0) * kernel_tt.mean() |
| ) |
| torch.testing.assert_close(loss, expected) |
| actual_gradient, = torch.autograd.grad( |
| loss, generated, retain_graph=True) |
| expected_gradient, = torch.autograd.grad(expected, generated) |
| torch.testing.assert_close( |
| actual_gradient, |
| expected_gradient, |
| rtol=2.0e-5, |
| atol=2.0e-6, |
| ) |
| self.assertEqual(int(metadata["global_samples"]), 4) |
| for partial_key in ( |
| "partial_student_reference_enabled", |
| "gradient_student_samples", |
| "student_reference_samples", |
| "student_reference_detached", |
| "student_repulsion_denominator", |
| "fixed_teacher_attraction_denominator", |
| "current_teacher_attraction_denominator", |
| "current_teacher_self_kernel_denominator", |
| "full_student_partial_gradient", |
| ): |
| self.assertNotIn(partial_key, metadata) |
|
|
| def test_partial_validation_rejects_ambiguous_or_stateful_modes(self): |
| generator = torch.Generator().manual_seed(9104) |
| gradient_students = torch.randn( |
| 2, 1, 2, generator=generator) |
| student_reference = torch.randn( |
| 4, 1, 2, generator=generator) |
| teacher = torch.randn(2, 1, 2, generator=generator) |
| objective = self._matched_objective() |
|
|
| with self.assertRaisesRegex(ValueError, "only valid"): |
| objective( |
| generated=gradient_students, |
| current_teacher=teacher, |
| gradient_student_indices=[0, 1], |
| ) |
| with self.assertRaisesRegex(TypeError, "must contain integers"): |
| objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| gradient_student_indices=[0.0, 1.0], |
| ) |
| with self.assertRaisesRegex(ValueError, "exactly one index"): |
| objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| gradient_student_indices=[0], |
| ) |
| with self.assertRaisesRegex(ValueError, "out-of-range"): |
| objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| gradient_student_indices=[0, 4], |
| ) |
| with self.assertRaisesRegex(ValueError, "must be unique"): |
| objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| gradient_student_indices=[0, 0], |
| ) |
| with self.assertRaisesRegex(ValueError, "must exactly match"): |
| objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| gradient_student_indices=[0, 1], |
| ) |
|
|
| with self.assertRaisesRegex(ValueError, "requires both"): |
| objective( |
| gradient_students=gradient_students, |
| current_teacher=teacher, |
| ) |
| with self.assertRaisesRegex(ValueError, "requires both"): |
| objective( |
| student_reference=student_reference, |
| current_teacher=teacher, |
| ) |
| with self.assertRaisesRegex(ValueError, "mutually exclusive"): |
| objective( |
| generated=gradient_students, |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=teacher, |
| ) |
| with self.assertRaisesRegex( |
| ValueError, "exactly match gradient_students" |
| ): |
| objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| current_teacher=torch.randn(3, 1, 2), |
| ) |
|
|
| fifo_objective = VideoNystromMMD( |
| self._pack(), repulsion_fifo_size=4) |
| with self.assertRaisesRegex(ValueError, "repulsion FIFO"): |
| fifo_objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| ) |
|
|
| ema_objective = VideoNystromMMD( |
| self._pack(), attraction_ema_enabled=True) |
| with self.assertRaisesRegex(ValueError, "attraction EMA"): |
| ema_objective( |
| gradient_students=gradient_students, |
| student_reference=student_reference, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
