def test_make_gauss_rectangular():
expected = torch.Tensor([
[0.496683, 0.182719, 0.024728, 0.001231, 0.000023],
[0.182719, 0.067219, 0.009097, 0.000453, 0.000008],
[0.024728, 0.009097, 0.001231, 0.000061, 0.000001],
])
actual = make_gauss(torch.Tensor([-1, -1]), [3, 5], sigma=1.0)
assert_allclose(actual, expected, rtol=0, atol=1e-5)
def test_rectangular(self):
expected = torch.Tensor([
[0.496683, 0.182719, 0.024728, 0.001231, 0.000023],
[0.182719, 0.067219, 0.009097, 0.000453, 0.000008],
[0.024728, 0.009097, 0.001231, 0.000061, 0.000001],
])
actual = make_gauss(torch.Tensor([-1, -1]), [3, 5], sigma=1.0)
self.assertEqual(expected, actual, 1e-5)
def test_make_gauss_2d():
expected = torch.Tensor([
[0.002969, 0.013306, 0.021938, 0.013306, 0.002969],
[0.013306, 0.059634, 0.098320, 0.059634, 0.013306],
[0.021938, 0.098320, 0.162103, 0.098320, 0.021938],
[0.013306, 0.059634, 0.098320, 0.059634, 0.013306],
[0.002969, 0.013306, 0.021938, 0.013306, 0.002969],
])
actual = make_gauss(torch.Tensor([0, 0]), [5, 5], sigma=1.0)
assert_allclose(actual, expected, rtol=0, atol=1e-5)
def test_2d(self):
expected = torch.Tensor([
[0.002969, 0.013306, 0.021938, 0.013306, 0.002969],
[0.013306, 0.059634, 0.098320, 0.059634, 0.013306],
[0.021938, 0.098320, 0.162103, 0.098320, 0.021938],
[0.013306, 0.059634, 0.098320, 0.059634, 0.013306],
[0.002969, 0.013306, 0.021938, 0.013306, 0.002969],
])
actual = make_gauss(torch.Tensor([0, 0]), [5, 5], sigma=1.0)
self.assertEqual(expected, actual, 1e-5)
def test_make_gauss_3d():
expected = torch.Tensor([[
[0.000035, 0.000002, 0.000000],
[0.009165, 0.000570, 0.000002],
[0.147403, 0.009165, 0.000035],
], [
[0.000142, 0.000009, 0.000000],
[0.036755, 0.002285, 0.000009],
[0.591145, 0.036755, 0.000142],
], [
[0.000035, 0.000002, 0.000000],
[0.009165, 0.000570, 0.000002],
[0.147403, 0.009165, 0.000035],
]])
actual = make_gauss(torch.Tensor([-1, 1, 0]), [3, 3, 3], sigma=0.6)
assert_allclose(actual, expected, rtol=0, atol=1e-5)
def test_3d(self):
expected = torch.Tensor([[
[0.000035, 0.000002, 0.000000],
[0.009165, 0.000570, 0.000002],
[0.147403, 0.009165, 0.000035],
], [
[0.000142, 0.000009, 0.000000],
[0.036755, 0.002285, 0.000009],
[0.591145, 0.036755, 0.000142],
], [
[0.000035, 0.000002, 0.000000],
[0.009165, 0.000570, 0.000002],
[0.147403, 0.009165, 0.000035],
]])
actual = make_gauss(torch.Tensor([-1, 1, 0]), [3, 3, 3], sigma=0.6)
self.assertEqual(expected, actual, 1e-5)
def calc_loss(mean, stddev):
hm = make_gauss(mean, [9, 9], sigma=stddev)
args = [hm]
if uses_mean: args.append(target_mean)
args.append(target_stddev)
return average_loss(loss_method(*args))
def test_make_gauss_unnormalized():
actual = make_gauss(torch.Tensor([0, 0]), [5, 5], sigma=1.0, normalize=False)
assert actual.max().item() == 1.0
def test_unnormalized(self):
actual = make_gauss(torch.Tensor([0, 0]), [5, 5],
sigma=1.0,
normalize=False)
self.assertEqual(1.0, actual.max().item())
for epoch in range(100):
total_loss = 0
total_correct = 0
model.train()
start = time.time()
for images, labels in train_loader: # Get Batch
cuda_images, cuda_labels = images.cuda(0), labels.cuda(1)
# Forward pass
coords, heatmaps = trained_model(cuda_images)
out = model(heatmaps.cuda(1))
loss = criterion(out, dsntnn.make_gauss(cuda_labels, (224, 224), 1.0))
total_loss += loss.item()
# Calculate gradients
optimizer.zero_grad()
loss.backward()
# Update model parameters with RMSprop
optimizer.step()
end = time.time()
torch.save(model.state_dict(),
checkpoints_path + "epoch_{}.pth".format(epoch))
writer.add_scalar('Training Loss', total_loss / len(train_loader), epoch)