31 lines
973 B
Python
31 lines
973 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class ProgressivePoolingLoss(nn.Module):
|
|
def __init__(self, initial_pool_size: int = 2, damping=1.8):
|
|
super(ProgressivePoolingLoss, self).__init__()
|
|
self._initial_pool_size = initial_pool_size
|
|
self._damping = damping
|
|
|
|
def forward(self, tensor_a, tensor_b):
|
|
assert (
|
|
tensor_a.size() == tensor_b.size()
|
|
), "Input tensors must have the same size."
|
|
|
|
max_pool_size = min(tensor_a.size(1), tensor_a.size(2), tensor_a.size(3))
|
|
|
|
loss = 0.0
|
|
damping = 1
|
|
|
|
for pool_size in range(self._initial_pool_size, max_pool_size):
|
|
pooled_a = F.avg_pool3d(tensor_a, pool_size) * (pool_size**3)
|
|
pooled_b = F.avg_pool3d(tensor_b, pool_size) * (pool_size**3)
|
|
|
|
diff = torch.square(pooled_a - pooled_b)
|
|
|
|
loss += diff.mean() / damping
|
|
damping *= self._damping
|
|
|
|
return loss
|