def collision_detection(self, obj):
if isinstance(obj, Object.Bullet):
pot = [
self.position[0], self.position[1] - self.earth,
self.position[2]
]
after_bullet = obj.position
before_bullet = obj.back()
before_to_after = util.Vec(after_bullet) - util.Vec(before_bullet)
pot_to_before = util.Vec(before_bullet) - util.Vec(pot)
pot_to_after = util.Vec(after_bullet) - util.Vec(pot)
if util.dot(before_to_after, -1 * pot_to_before) < 0:
if util.norm(pot_to_before) < self.radius:
return True
else:
return False
elif util.dot(before_to_after, pot_to_after) < 0:
if util.norm(pot_to_after) < self.radius:
return True
else:
return False
else:
if util.norm(util.cross3d(before_to_after, pot_to_before)
) / util.norm(before_to_after) < self.radius:
return True
else:
return False
def extra_checks(A, X, B):
x = u.Vec(X)
u.check_equal(u.Vec(A @ X @ B), x @ u.Kron(B, A.t()))
u.check_equal(u.Vec(A @ X @ B), u.Kron(B.t(), A) @ x)
u.check_equal(u.Vecr(A @ X @ B), u.Kron(A, B.t()) @ u.Vecr(X))
u.check_equal(u.Vecr(A @ X @ B), u.Vecr(X) @ u.Kron(A.t(), B))
u.check_equal(u.Vecr(A @ X @ B),
u.Vecr(X) @ u.Kron(A.t(), B).normal_form())
u.check_equal(u.Vecr(A @ X @ B),
u.matmul(u.Kron(A, B.t()).normal_form(), u.Vecr(X)))
u.check_equal(u.Vec(A @ X @ B),
u.matmul(u.Kron(B.t(), A).normal_form(), x))
u.check_equal(u.Vec(A @ X @ B), x @ u.Kron(B, A.t()).normal_form())
u.check_equal(u.Vec(A @ X @ B),
x.normal_form() @ u.Kron(B, A.t()).normal_form())
u.check_equal(u.Vec(A @ X @ B),
u.Kron(B.t(), A).normal_form() @ x.normal_form())
u.check_equal(u.Vecr(A @ X @ B),
u.Kron(A, B.t()).normal_form() @ u.Vecr(X).normal_form())
u.check_equal(u.Vecr(A @ X @ B),
u.Vecr(X).normal_form() @ u.Kron(A.t(), B).normal_form())
def test_explicit_hessian():
"""Check computation of hessian of loss(B'WA) from https://github.com/yaroslavvb/kfac_pytorch/blob/master/derivation.pdf
"""
torch.set_default_dtype(torch.float64)
A = torch.tensor([[-1., 4], [3, 0]])
B = torch.tensor([[-4., 3], [2, 6]])
X = torch.tensor([[-5., 0], [-2, -6]], requires_grad=True)
Y = B.t() @ X @ A
u.check_equal(Y, [[-52, 64], [-81, -108]])
loss = torch.sum(Y * Y) / 2
hess0 = u.hessian(loss, X).reshape([4, 4])
hess1 = u.Kron(A @ A.t(), B @ B.t())
u.check_equal(loss, 12512.5)
# PyTorch autograd computes Hessian with respect to row-vectorized parameters, whereas
# autograd_lib uses math convention and does column-vectorized.
# Commuting order of Kronecker product switches between two representations
u.check_equal(hess1.commute(), hess0)
# Do a test using Linear layers instead of matrix multiplies
model: u.SimpleFullyConnected2 = u.SimpleFullyConnected2([2, 2, 2], bias=False)
model.layers[0].weight.data.copy_(X)
# Transpose to match previous results, layers treat dim0 as batch dimension
u.check_equal(model.layers[0](A.t()).t(), [[5, -20], [-16, -8]]) # XA = (A'X0)'
model.layers[1].weight.data.copy_(B.t())
u.check_equal(model(A.t()).t(), Y)
Y = model(A.t()).t() # transpose to data-dimension=columns
loss = torch.sum(Y * Y) / 2
loss.backward()
u.check_equal(model.layers[0].weight.grad, [[-2285, -105], [-1490, -1770]])
G = B @ Y @ A.t()
u.check_equal(model.layers[0].weight.grad, G)
u.check_equal(hess0, u.Kron(B @ B.t(), A @ A.t()))
# compute newton step
u.check_equal(u.Kron([email protected](), [email protected]()).pinv() @ u.vec(G), u.v2c([-5, -2, 0, -6]))
# compute Newton step using factored representation
autograd_lib.add_hooks(model)
Y = model(A.t())
n = 2
loss = torch.sum(Y * Y) / 2
autograd_lib.backprop_hess(Y, hess_type='LeastSquares')
autograd_lib.compute_hess(model, method='kron', attr_name='hess_kron', vecr_order=False, loss_aggregation='sum')
param = model.layers[0].weight
hess2 = param.hess_kron
print(hess2)
u.check_equal(hess2, [[425, 170, -75, -30], [170, 680, -30, -120], [-75, -30, 225, 90], [-30, -120, 90, 360]])
# Gradient test
model.zero_grad()
loss.backward()
u.check_close(u.vec(G).flatten(), u.Vec(param.grad))
# Newton step test
# Method 0: PyTorch native autograd
newton_step0 = param.grad.flatten() @ torch.pinverse(hess0)
newton_step0 = newton_step0.reshape(param.shape)
u.check_equal(newton_step0, [[-5, 0], [-2, -6]])
# Method 1: colummn major order
ihess2 = hess2.pinv()
u.check_equal(ihess2.LL, [[1/16, 1/48], [1/48, 17/144]])
u.check_equal(ihess2.RR, [[2/45, -(1/90)], [-(1/90), 1/36]])
u.check_equal(torch.flatten(hess2.pinv() @ u.vec(G)), [-5, -2, 0, -6])
newton_step1 = (ihess2 @ u.Vec(param.grad)).matrix_form()
# Method2: row major order
ihess2_rowmajor = ihess2.commute()
newton_step2 = ihess2_rowmajor @ u.Vecr(param.grad)
newton_step2 = newton_step2.matrix_form()
u.check_equal(newton_step0, newton_step1)
u.check_equal(newton_step0, newton_step2)
def _test_explicit_hessian_refactored():
"""Check computation of hessian of loss(B'WA) from https://github.com/yaroslavvb/kfac_pytorch/blob/master/derivation.pdf
"""
torch.set_default_dtype(torch.float64)
A = torch.tensor([[-1., 4], [3, 0]])
B = torch.tensor([[-4., 3], [2, 6]])
X = torch.tensor([[-5., 0], [-2, -6]], requires_grad=True)
Y = B.t() @ X @ A
u.check_equal(Y, [[-52, 64], [-81, -108]])
loss = torch.sum(Y * Y) / 2
hess0 = u.hessian(loss, X).reshape([4, 4])
hess1 = u.Kron(A @ A.t(), B @ B.t())
u.check_equal(loss, 12512.5)
# Do a test using Linear layers instead of matrix multiplies
model: u.SimpleFullyConnected2 = u.SimpleFullyConnected2([2, 2, 2], bias=False)
model.layers[0].weight.data.copy_(X)
# Transpose to match previous results, layers treat dim0 as batch dimension
u.check_equal(model.layers[0](A.t()).t(), [[5, -20], [-16, -8]]) # XA = (A'X0)'
model.layers[1].weight.data.copy_(B.t())
u.check_equal(model(A.t()).t(), Y)
Y = model(A.t()).t() # transpose to data-dimension=columns
loss = torch.sum(Y * Y) / 2
loss.backward()
u.check_equal(model.layers[0].weight.grad, [[-2285, -105], [-1490, -1770]])
G = B @ Y @ A.t()
u.check_equal(model.layers[0].weight.grad, G)
autograd_lib.register(model)
activations_dict = autograd_lib.ModuleDict() # todo(y): make save_activations ctx manager automatically create A
with autograd_lib.save_activations(activations_dict):
Y = model(A.t())
Acov = autograd_lib.ModuleDict(autograd_lib.SecondOrderCov)
for layer, activations in activations_dict.items():
print(layer, activations)
Acov[layer].accumulate(activations, activations)
autograd_lib.set_default_activations(activations_dict)
autograd_lib.set_default_Acov(Acov)
B = autograd_lib.ModuleDict(autograd_lib.SymmetricFourthOrderCov)
autograd_lib.backward_accum(Y, "identity", B, retain_graph=False)
print(B[model.layers[0]])
autograd_lib.backprop_hess(Y, hess_type='LeastSquares')
autograd_lib.compute_hess(model, method='kron', attr_name='hess_kron', vecr_order=False, loss_aggregation='sum')
param = model.layers[0].weight
hess2 = param.hess_kron
print(hess2)
u.check_equal(hess2, [[425, 170, -75, -30], [170, 680, -30, -120], [-75, -30, 225, 90], [-30, -120, 90, 360]])
# Gradient test
model.zero_grad()
loss.backward()
u.check_close(u.vec(G).flatten(), u.Vec(param.grad))
# Newton step test
# Method 0: PyTorch native autograd
newton_step0 = param.grad.flatten() @ torch.pinverse(hess0)
newton_step0 = newton_step0.reshape(param.shape)
u.check_equal(newton_step0, [[-5, 0], [-2, -6]])
# Method 1: colummn major order
ihess2 = hess2.pinv()
u.check_equal(ihess2.LL, [[1/16, 1/48], [1/48, 17/144]])
u.check_equal(ihess2.RR, [[2/45, -(1/90)], [-(1/90), 1/36]])
u.check_equal(torch.flatten(hess2.pinv() @ u.vec(G)), [-5, -2, 0, -6])
newton_step1 = (ihess2 @ u.Vec(param.grad)).matrix_form()
# Method2: row major order
ihess2_rowmajor = ihess2.commute()
newton_step2 = ihess2_rowmajor @ u.Vecr(param.grad)
newton_step2 = newton_step2.matrix_form()
u.check_equal(newton_step0, newton_step1)
u.check_equal(newton_step0, newton_step2)
def test_kron():
"""Test kron, vec and vecr identities"""
torch.set_default_dtype(torch.float64)
a = torch.tensor([1, 2, 3, 4]).reshape(2, 2)
b = torch.tensor([5, 6, 7, 8]).reshape(2, 2)
u.check_close(u.Kron(a, b).trace(), 65)
a = torch.tensor([[2., 7, 9], [1, 9, 8], [2, 7, 5]])
b = torch.tensor([[6., 6, 1], [10, 7, 7], [7, 10, 10]])
Ck = u.Kron(a, b)
u.check_close(a.flatten().norm() * b.flatten().norm(), Ck.frobenius_norm())
u.check_close(Ck.frobenius_norm(), 4 * math.sqrt(11635.))
Ci = [[
0, 5 / 102, -(7 / 204), 0, -(70 / 561), 49 / 561, 0, 125 / 1122,
-(175 / 2244)
],
[
1 / 20, -(53 / 1020), 8 / 255, -(7 / 55), 371 / 2805,
-(224 / 2805), 5 / 44, -(265 / 2244), 40 / 561
],
[
-(1 / 20), 3 / 170, 3 / 170, 7 / 55, -(42 / 935), -(42 / 935),
-(5 / 44), 15 / 374, 15 / 374
],
[
0, -(5 / 102), 7 / 204, 0, 20 / 561, -(14 / 561), 0, 35 / 1122,
-(49 / 2244)
],
[
-(1 / 20), 53 / 1020, -(8 / 255), 2 / 55, -(106 / 2805),
64 / 2805, 7 / 220, -(371 / 11220), 56 / 2805
],
[
1 / 20, -(3 / 170), -(3 / 170), -(2 / 55), 12 / 935, 12 / 935,
-(7 / 220), 21 / 1870, 21 / 1870
], [0, 5 / 102, -(7 / 204), 0, 0, 0, 0, -(5 / 102), 7 / 204],
[
1 / 20, -(53 / 1020), 8 / 255, 0, 0, 0, -(1 / 20), 53 / 1020,
-(8 / 255)
],
[
-(1 / 20), 3 / 170, 3 / 170, 0, 0, 0, 1 / 20, -(3 / 170),
-(3 / 170)
]]
C = Ck.expand()
C0 = u.to_numpy(C)
Ci = torch.tensor(Ci)
u.check_close(C @ Ci @ C, C)
u.check_close(Ck.inv().expand(), torch.inverse(Ck.expand()))
u.check_close(Ck.inv().expand_vec(), torch.inverse(Ck.expand_vec()))
u.check_close(Ck.pinv().expand(), torch.pinverse(Ck.expand()))
u.check_close(linalg.pinv(C0), Ci, rtol=1e-5, atol=1e-6)
u.check_close(torch.pinverse(C), Ci, rtol=1e-5, atol=1e-6)
u.check_close(Ck.inv().expand(), Ci, rtol=1e-5, atol=1e-6)
u.check_close(Ck.pinv().expand(), Ci, rtol=1e-5, atol=1e-6)
Ck2 = u.Kron(b, 2 * a)
u.check_close((Ck @ Ck2).expand(), Ck.expand() @ Ck2.expand())
u.check_close((Ck @ Ck2).expand_vec(), Ck.expand_vec() @ Ck2.expand_vec())
d2 = 3
d1 = 2
G = torch.randn(d2, d1)
g = u.vec(G)
H = u.Kron(u.random_cov(d1), u.random_cov(d2))
Gt = G.t()
gt = g.reshape(1, -1)
vecX = u.Vec([1, 2, 3, 4], shape=(2, 2))
K = u.Kron([[5, 6], [7, 8]], [[9, 10], [11, 12]])
u.check_equal(vecX @ K, [644, 706, 748, 820])
u.check_equal(K @ vecX, [543, 655, 737, 889])
u.check_equal(u.matmul(vecX @ K, vecX), 7538)
u.check_equal(vecX @ (vecX @ K), 7538)
u.check_equal(vecX @ vecX, 30)
vecX = u.Vec([1, 2], shape=(1, 2))
K = u.Kron([[5]], [[9, 10], [11, 12]])
u.check_equal(vecX.norm()**2, 5)
# check kronecker rules
X = torch.tensor([[1., 2], [3, 4]])
A = torch.tensor([[5., 6], [7, 8]])
B = torch.tensor([[9., 10], [11, 12]])
x = u.Vec(X)
# kron/vec/vecr identities
u.check_equal(u.Vec(A @ X @ B), x @ u.Kron(B, A.t()))
u.check_equal(u.Vec(A @ X @ B), u.Kron(B.t(), A) @ x)
u.check_equal(u.Vecr(A @ X @ B), u.Kron(A, B.t()) @ u.Vecr(X))
u.check_equal(u.Vecr(A @ X @ B), u.Vecr(X) @ u.Kron(A.t(), B))
def extra_checks(A, X, B):
x = u.Vec(X)
u.check_equal(u.Vec(A @ X @ B), x @ u.Kron(B, A.t()))
u.check_equal(u.Vec(A @ X @ B), u.Kron(B.t(), A) @ x)
u.check_equal(u.Vecr(A @ X @ B), u.Kron(A, B.t()) @ u.Vecr(X))
u.check_equal(u.Vecr(A @ X @ B), u.Vecr(X) @ u.Kron(A.t(), B))
u.check_equal(u.Vecr(A @ X @ B),
u.Vecr(X) @ u.Kron(A.t(), B).normal_form())
u.check_equal(u.Vecr(A @ X @ B),
u.matmul(u.Kron(A, B.t()).normal_form(), u.Vecr(X)))
u.check_equal(u.Vec(A @ X @ B),
u.matmul(u.Kron(B.t(), A).normal_form(), x))
u.check_equal(u.Vec(A @ X @ B), x @ u.Kron(B, A.t()).normal_form())
u.check_equal(u.Vec(A @ X @ B),
x.normal_form() @ u.Kron(B, A.t()).normal_form())
u.check_equal(u.Vec(A @ X @ B),
u.Kron(B.t(), A).normal_form() @ x.normal_form())
u.check_equal(u.Vecr(A @ X @ B),
u.Kron(A, B.t()).normal_form() @ u.Vecr(X).normal_form())
u.check_equal(u.Vecr(A @ X @ B),
u.Vecr(X).normal_form() @ u.Kron(A.t(), B).normal_form())
# shape checks
d1, d2 = 3, 4
extra_checks(torch.ones((d1, d1)), torch.ones((d1, d2)),
torch.ones((d2, d2)))
A = torch.rand(d1, d1)
B = torch.rand(d2, d2)
#x = torch.rand((d1*d2))
#X = x.t().reshape(d1, d2)
# X = torch.rand((d1, d2))
# x = u.vec(X)
x = torch.rand((d1 * d2))