File size: 1,856 Bytes
36c95ba |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import logging
import torch
import torch.nn as nn
import torch.optim as optim
import kornia
from kornia.testing import assert_close
logger = logging.getLogger(__name__)
class TestIntegrationSoftArgmax2d:
# optimization
lr = 1e-3
num_iterations = 500
# data params
height = 240
width = 320
def generate_sample(self, base_target, std_val=1.0):
"""Generate a random sample around the given point.
The standard deviation is in pixel.
"""
noise = std_val * torch.rand_like(base_target)
return base_target + noise
def test_regression_2d(self, device):
# create the parameters to estimate: the heatmap
params = nn.Parameter(torch.rand(1, 1, self.height, self.width).to(device))
# generate base sample
target = torch.zeros(1, 1, 2).to(device)
target[..., 0] = self.width / 2
target[..., 1] = self.height / 2
# create the optimizer and pass the heatmap
optimizer = optim.Adam([params], lr=self.lr)
# loss criterion
criterion = nn.MSELoss()
# spatial soft-argmax2d module
soft_argmax2d = kornia.geometry.SpatialSoftArgmax2d(normalized_coordinates=False)
# NOTE: check where this comes from
temperature = (self.height * self.width) ** (0.5)
for _ in range(self.num_iterations):
x = params
sample = self.generate_sample(target).to(device)
pred = soft_argmax2d(temperature * x)
loss = criterion(pred, sample)
logger.debug(f"Loss: {loss.item():.3f} Pred: {pred}")
optimizer.zero_grad()
loss.backward()
optimizer.step()
assert_close(pred[..., 0], target[..., 0], rtol=1e-2, atol=1e-2)
assert_close(pred[..., 1], target[..., 1], rtol=1e-2, atol=1e-2)
|