def test_jit(self, device, dtype):
B, C, H, W = 2, 1, 13, 13
patches = torch.ones(B, C, H, W, device=device, dtype=dtype)
tfeat = PatchAffineShapeEstimator(W).to(patches.device,
patches.dtype).eval()
tfeat_jit = torch.jit.script(
PatchAffineShapeEstimator(W).to(patches.device,
patches.dtype).eval())
assert_close(tfeat_jit(patches), tfeat(patches))
def test_toy(self, device):
aff = PatchAffineShapeEstimator(19).to(device)
inp = torch.zeros(1, 1, 19, 19, device=device)
inp[:, :, 5:-5, 1:-1] = 1
abc = aff(inp)
expected = torch.tensor([[[0.4146, 0.0000, 1.0000]]], device=device)
assert_close(abc, expected, atol=1e-4, rtol=1e-4)
def test_gradcheck(self, device):
batch_size, channels, height, width = 1, 1, 13, 13
ori = PatchAffineShapeEstimator(width).to(device)
patches = torch.rand(batch_size,
channels,
height,
width,
device=device)
patches = utils.tensor_to_gradcheck_var(patches) # to var
assert gradcheck(ori, (patches, ),
raise_exception=True,
nondet_tol=1e-4)
def test_print(self, device):
sift = PatchAffineShapeEstimator(32)
sift.__repr__()
def test_shape_batch(self, device):
inp = torch.rand(2, 1, 32, 32, device=device)
ori = PatchAffineShapeEstimator(32).to(device)
ang = ori(inp)
assert ang.shape == torch.Size([2, 1, 3])