File size: 1,353 Bytes
3f69d68 | 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 | import torch
from lightnest.global_patch_context import GlobalPatchNetwork,freeze,examples
from lightnest.learned_patch import batch,PatchNetwork
from lightnest.probe_receptive_field import measure
def test_frozen_pairs_and_labels():
pairs=freeze(12,2,'combined');ex=examples(pairs)
assert len(ex)==4
assert all(sum(a!=b for a,b in zip(e.binary,e.target))==1 for e in ex)
assert freeze(12,2,'combined')==pairs
def test_global_context_breaks_local_invariance():
torch.set_num_threads(2);torch.manual_seed(82)
pairs=freeze(12,2,'combined');local=PatchNetwork();global_model=GlobalPatchNetwork()
global_model.load_state_dict(local.state_dict(),strict=False)
with torch.no_grad(): global_model.global_scale.fill_(1.)
assert all(x['pair_difference_change']==0 for x in measure(local,pairs))
assert any(x['pair_difference_change']>1e-5 for x in measure(global_model,pairs))
assert not global_model.native_packer_compatible
def test_outputs_and_gradients():
torch.set_num_threads(2);model=GlobalPatchNetwork();c,b,p,y=batch(examples(freeze(13,1,'train')))
scores,bits=model(c,b)
assert scores.shape==(2,257) and bits.shape==(2,256,8)
torch.nn.functional.cross_entropy(scores,p).backward()
assert model.global_attention.in_proj_weight.grad.abs().sum()>0
assert torch.isfinite(scores).all()
|