File size: 2,460 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 29 30 31 32 33 34 35 36 37 38 39 | import copy,random,unittest
import torch
from torch.nn import functional as F
from lightnest.gradient_interference_experiment import RoutedXorPolicy,freeze,teacher_metrics
from lightnest.xor_patch_experiment import inputs,XorPolicy
from lightnest.multifamily_patch import make_task,teacher_states
class GradientRoutingTests(unittest.TestCase):
def setUp(self):
torch.set_num_threads(2);torch.manual_seed(17)
self.states=teacher_states(make_task(random.Random(31)))
def test_byte_gradient_isolation(self):
for detached in [False,True]:
m=RoutedXorPolicy(detached);c,b,p,y=inputs(self.states);mask=p<256
_,bits=m(c,b);F.binary_cross_entropy_with_logits(bits[torch.arange(len(p))[mask],p[mask]],y[mask]).backward()
backbone=sum(float(v.grad.abs().sum()) for k,v in m.named_parameters() if not k.startswith('replacement.') and v.grad is not None)
self.assertGreater(float(m.replacement.weight.grad.abs().sum()),0)
if detached:self.assertEqual(backbone,0)
else:self.assertGreater(backbone,0)
def test_pointer_gradient_unaffected(self):
a=RoutedXorPolicy();b=copy.deepcopy(a);b.detach_byte=True;c,x,p,y=inputs(self.states)
for m in [a,b]:F.cross_entropy(m(c,x)[0],p).backward()
self.assertGreater(float(a.fusion.weight.grad.abs().sum()),0)
for (k,v),(l,w) in zip(a.named_parameters(),b.named_parameters()):
if v.grad is None:self.assertIsNone(w.grad)
else:self.assertTrue(torch.equal(v.grad,w.grad),k)
def test_inference_equivalent_to_existing_xor(self):
a=RoutedXorPolicy();b=XorPolicy(copy.deepcopy(a));c,x,_,_=inputs(self.states)
# Existing wrapper expects absolute-network logits, so compare with plain PatchNetwork.
from lightnest.learned_patch import PatchNetwork
net=PatchNetwork();net.load_state_dict(a.state_dict());b=XorPolicy(net)
for lhs,rhs in zip(a(c,x),b(c,x)):self.assertTrue(torch.equal(lhs,rhs))
a.detach_byte=True
for lhs,rhs in zip(a(c,x),b(c,x)):self.assertTrue(torch.equal(lhs,rhs))
def test_frozen_disjoint_and_reproducible(self):
train,suites=freeze();train2,suites2=freeze();self.assertEqual(train,train2);self.assertEqual(suites,suites2)
self.assertFalse({(t.command,t.binary) for t in train}&{(t.command,t.binary) for ts in suites.values() for t in ts})
if __name__=='__main__':unittest.main()
|