def test_TotalVariationLoss(input_image):
loss = paper.TotalVariationLoss()
actual = loss(input_image)
desired = F.total_variation_loss(input_image, reduction="sum")
assert actual == ptu.approx(desired)
def test_call(self):
torch.manual_seed(0)
image = torch.rand(1, 3, 128, 128)
exponent = 3.0
op = loss.TotalVariationLoss(exponent=exponent)
actual = op(image)
desired = F.total_variation_loss(image, exponent=exponent)
ptu.assert_allclose(actual, desired)
def test_total_variation_loss():
def get_checkerboard(size):
return ((torch.arange(size**2).view(size, size) +
torch.arange(size).view(size, 1)) % 2).bool()
size = 128
checkerboard = get_checkerboard(size)
input = checkerboard.float().view(1, 1, size, size).repeat(1, 3, 1, 1)
actual = F.total_variation_loss(input)
desired = 2.0
assert desired == ptu.approx(actual)
def test_TotalVariationLoss(subtests, input_image):
configs = ((True, 1.0 / 2.0), (False, 1.0))
for impl_params, score_correction_factor in configs:
with subtests.test(impl_params=impl_params):
loss = paper.TotalVariationLoss(impl_params=impl_params, )
actual = loss(input_image)
score = F.total_variation_loss(input_image,
exponent=loss.exponent,
reduction="sum")
desired = score * score_correction_factor
assert actual == ptu.approx(desired)
def calculate_score(self, input_repr: torch.Tensor) -> torch.Tensor:
score = F.total_variation_loss(input_repr,
exponent=self.exponent,
reduction=self.loss_reduction)
return score * self.score_correction_factor
def calculate_score(self, input_repr: torch.Tensor) -> torch.Tensor:
return F.total_variation_loss(input_repr,
exponent=self.exponent,
reduction=self.loss_reduction)