Buckets:
| import unittest | |
| import torch | |
| from opd.engine import apply_optimizer_update | |
| from shared.ema import EMATeacher | |
| class OPDEngineTest(unittest.TestCase): | |
| def test_zero_accepted_update_advances_budget_scheduler_and_ema_without_parameter_change(self): | |
| model = torch.nn.Linear(1, 1, bias=False) | |
| before = model.weight.detach().clone() | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0) | |
| teacher = EMATeacher(model, decay=0.9) | |
| result = apply_optimizer_update( | |
| losses=[], | |
| parameters=list(model.parameters()), | |
| model=model, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ema_teacher=teacher, | |
| max_grad_norm=5.0, | |
| ) | |
| self.assertTrue(result.successful) | |
| self.assertTrue(result.zero_gradient) | |
| self.assertEqual(scheduler.last_epoch, 1) | |
| self.assertEqual(teacher.update_count, 1) | |
| torch.testing.assert_close(model.weight, before) | |
| def test_finite_loss_updates_student_then_ema(self): | |
| model = torch.nn.Linear(1, 1, bias=False) | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0) | |
| teacher = EMATeacher(model, decay=0.5) | |
| before = model.weight.detach().clone() | |
| loss = model(torch.ones(1, 1)).square().sum() | |
| result = apply_optimizer_update( | |
| losses=[loss], | |
| parameters=list(model.parameters()), | |
| model=model, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ema_teacher=teacher, | |
| max_grad_norm=5.0, | |
| ) | |
| self.assertTrue(result.successful) | |
| self.assertFalse(result.zero_gradient) | |
| self.assertFalse(torch.equal(model.weight, before)) | |
| self.assertEqual(teacher.update_count, 1) | |
| def test_distributed_sum_happens_before_global_clipping_and_step(self): | |
| model = torch.nn.Linear(1, 1, bias=False) | |
| model.weight.data.fill_(1.0) | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _step: 1.0) | |
| finite_values = [] | |
| def synchronize(parameters): | |
| parameters[0].grad.mul_(2.0) | |
| result = apply_optimizer_update( | |
| losses=[model.weight.square().sum()], | |
| parameters=list(model.parameters()), | |
| model=model, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ema_teacher=None, | |
| max_grad_norm=100.0, | |
| synchronize_gradients=synchronize, | |
| synchronize_finite=lambda value: finite_values.append(value) or value, | |
| ) | |
| self.assertTrue(result.successful) | |
| self.assertEqual(finite_values, [True]) | |
| self.assertAlmostEqual(float(model.weight), 0.6, places=6) | |
| def test_nonfinite_post_step_state_aborts_before_scheduler_and_ema(self): | |
| model = torch.nn.Linear(1, 1, bias=False) | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=0.1) | |
| scheduler = torch.optim.lr_scheduler.LambdaLR( | |
| optimizer, lambda _step: 1.0 | |
| ) | |
| teacher = EMATeacher(model, decay=0.9) | |
| loss = model(torch.ones(1, 1)).square().sum() | |
| original_step = optimizer.step | |
| def corrupting_step(): | |
| original_step() | |
| model.weight.data.fill_(float("inf")) | |
| optimizer.step = corrupting_step | |
| with self.assertRaisesRegex(FloatingPointError, "optimizer state"): | |
| apply_optimizer_update( | |
| losses=[loss], | |
| parameters=list(model.parameters()), | |
| model=model, | |
| optimizer=optimizer, | |
| scheduler=scheduler, | |
| ema_teacher=teacher, | |
| max_grad_norm=5.0, | |
| ) | |
| self.assertEqual(scheduler.last_epoch, 0) | |
| self.assertEqual(teacher.update_count, 0) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 4.1 kB
- Xet hash:
- ae3b0bcf9ea82d31605e120c6f0de1872ed77d8356203c68bf1f140772b73b4e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.