| """Capacity, inherited behavior, and trainable new features.""" |
| import random |
|
|
| import pytest |
| import torch |
|
|
| from lightnest.learned_patch import PatchNetwork, batch, load_predictor, make_example |
| from lightnest.patch_capacity import parameter_counts, widen_model |
|
|
|
|
| @pytest.mark.parametrize('width,total,heads', [ |
| (1, 107994, 714), (2, 405162, 1418), (3, 891514, 2122), |
| (4, 1567050, 2826), (8, 6161034, 5642), |
| ]) |
| def test_declared_capacity_round_trip(tmp_path, width, total, heads): |
| torch.set_num_threads(2) |
| model = PatchNetwork(width=width) |
| counts = parameter_counts(model) |
| assert counts['total'] == counts['trainable'] == total |
| assert counts['head'] == heads |
| path = tmp_path / 'model.pt' |
| torch.save({'state_dict': model.state_dict(), 'threshold': .55, |
| 'model_config': model.model_config}, path) |
| restored = load_predictor(path).model |
| assert restored.width == width |
| assert all(torch.equal(restored.state_dict()[k], value) |
| for k, value in model.state_dict().items()) |
|
|
|
|
| @pytest.mark.parametrize('before,after', [(1, 1), (1, 2), (1, 4), (1, 8), (2, 5)]) |
| def test_widening_preserves_outputs_and_parent(before, after): |
| torch.set_num_threads(2) |
| torch.manual_seed(96171) |
| model = PatchNetwork(width=before).eval() |
| old = {k: v.clone() for k, v in model.state_dict().items()} |
| commands, binaries, _, _ = batch([ |
| make_example(random.Random(10291 + i), 'combined') for i in range(8) |
| ]) |
| random_before = torch.random.get_rng_state().clone() |
| expanded = widen_model(model, after, seed=21419) |
| assert torch.equal(torch.random.get_rng_state(), random_before) |
| with torch.no_grad(): |
| for expected, actual in zip(model(commands, binaries), expanded(commands, binaries)): |
| torch.testing.assert_close(actual, expected, atol=3e-6, rtol=1e-6) |
| assert all(torch.equal(model.state_dict()[k], v) for k, v in old.items()) |
| assert next(model.parameters()).data_ptr() != next(expanded.parameters()).data_ptr() |
|
|
|
|
| def test_new_features_receive_updates_after_connections_learn(): |
| torch.set_num_threads(2) |
| torch.manual_seed(19019) |
| expanded = widen_model(PatchNetwork(), 2, seed=72419) |
| before = {k: v.clone() for k, v in expanded.state_dict().items()} |
| commands, binaries, pointers, replacements = batch([ |
| make_example(random.Random(62291 + i)) for i in range(8) |
| ]) |
| optimizer = torch.optim.AdamW(expanded.parameters(), lr=.003, weight_decay=0) |
| for _ in range(4): |
| logits, bits = expanded(commands, binaries) |
| positive = pointers != 256 |
| loss = torch.nn.functional.cross_entropy(logits, pointers) |
| loss += torch.nn.functional.binary_cross_entropy_with_logits( |
| bits[positive, pointers[positive]], replacements[positive]) |
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
| |
| assert not torch.equal(expanded.command_embedding.weight[:, 24:], |
| before['command_embedding.weight'][:, 24:]) |
| assert not torch.equal(expanded.binary_embedding.weight[:, 24:], |
| before['binary_embedding.weight'][:, 24:]) |
| assert not torch.equal(expanded.binary_conv[2].weight[64:], |
| before['binary_conv.2.weight'][64:]) |
| assert torch.count_nonzero(expanded.pointer.weight[:, 64:]) |
|
|
|
|
| @pytest.mark.parametrize('config', [None, {'width': True}, {'width': 0}, |
| {'width': 9}, {'width': 2.0}, {'depth': 2}]) |
| def test_invalid_or_undeclared_dimensions_reject(tmp_path, config): |
| path = tmp_path / 'bad.pt' |
| torch.save({'state_dict': PatchNetwork().state_dict(), 'threshold': .55, |
| 'model_config': config}, path) |
| with pytest.raises(ValueError): |
| load_predictor(path) |
|
|
|
|
| def test_wider_tensors_require_matching_metadata(tmp_path): |
| path = tmp_path / 'bad.pt' |
| torch.save({'state_dict': PatchNetwork(width=2).state_dict(), 'threshold': .55}, path) |
| with pytest.raises(ValueError, match='checkpoint tensor'): |
| load_predictor(path) |
|
|
|
|
| def test_checkpoint_size_limit_precedes_deserialization(tmp_path, monkeypatch): |
| path = tmp_path / 'oversized.pt' |
| with path.open('wb') as stream: |
| stream.truncate(32 * 1024**2 + 1) |
| def unexpected_load(*args, **kwargs): |
| raise AssertionError('oversized file must not reach torch.load') |
| monkeypatch.setattr(torch, 'load', unexpected_load) |
| with pytest.raises(ValueError, match='32 MiB'): |
| load_predictor(path) |
|
|
