112 lines
4.1 KiB
Python
112 lines
4.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
def gaussian_radial_basis_function(x, mu, sigma):
|
|
# (kernel_size) -> (batch_size, kernel_size, c*h*w)
|
|
mu = mu.view(1, mu.size(0), 1).expand(x.size(0), -1, x.size(1) * x.size(2) * x.size(3))
|
|
mu = mu.to(x.device)
|
|
# (batch_size, c, h, w) -> (batch_size, kernel_size, c*h*w)
|
|
x = x.view(x.size(0), 1, -1).expand(-1, mu.size(1), -1)
|
|
return torch.exp((x - mu).pow(2) / (2 * sigma ** 2))
|
|
|
|
|
|
class MyLoss(torch.nn.Module):
|
|
def __init__(self):
|
|
super(MyLoss, self).__init__()
|
|
|
|
def forward(self, fakeI, realI):
|
|
def batch_ERSMI(I1, I2):
|
|
batch_size = I1.shape[0]
|
|
img_size = I1.shape[1] * I1.shape[2] * I1.shape[3]
|
|
if I2.shape[1] == 1 and I1.shape[1] != 1:
|
|
I2 = I2.repeat(1, 3, 1, 1)
|
|
|
|
def kernel_F(y, mu_list, sigma):
|
|
tmp_mu = mu_list.view(-1, 1).repeat(1, img_size).repeat(batch_size, 1, 1).cuda() # [81, 784]
|
|
tmp_y = y.view(batch_size, 1, -1).repeat(1, 81, 1)
|
|
tmp_y = tmp_mu - tmp_y
|
|
mat_L = torch.exp(tmp_y.pow(2) / (2 * sigma ** 2))
|
|
return mat_L
|
|
|
|
mu = torch.Tensor([-1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75, 1.0]).cuda()
|
|
|
|
x_mu_list = mu.repeat(9).view(-1, 81)
|
|
y_mu_list = mu.unsqueeze(0).t().repeat(1, 9).view(-1, 81)
|
|
|
|
mat_K = kernel_F(I1, x_mu_list, 1)
|
|
mat_L = kernel_F(I2, y_mu_list, 1)
|
|
|
|
H1 = ((mat_K.matmul(mat_K.transpose(1, 2))).mul(mat_L.matmul(mat_L.transpose(1, 2))) / (
|
|
img_size ** 2)).cuda()
|
|
H2 = ((mat_K.mul(mat_L)).matmul((mat_K.mul(mat_L)).transpose(1, 2)) / img_size).cuda()
|
|
h2 = ((mat_K.sum(2).view(batch_size, -1, 1)).mul(mat_L.sum(2).view(batch_size, -1, 1)) / (
|
|
img_size ** 2)).cuda()
|
|
H2 = 0.5 * H1 + 0.5 * H2
|
|
tmp = H2 + 0.05 * torch.eye(len(H2[0])).cuda()
|
|
alpha = (tmp.inverse())
|
|
|
|
alpha = alpha.matmul(h2)
|
|
ersmi = (2 * (alpha.transpose(1, 2)).matmul(h2) - ((alpha.transpose(1, 2)).matmul(H2)).matmul(
|
|
alpha) - 1).squeeze()
|
|
ersmi = -ersmi.mean()
|
|
return ersmi
|
|
|
|
batch_loss = batch_ERSMI(fakeI, realI)
|
|
return batch_loss
|
|
|
|
|
|
class MGCLoss(nn.Module):
|
|
"""
|
|
Minimal Geometry-Distortion Constraint Loss from https://openreview.net/forum?id=R5M7Mxl1xZ
|
|
"""
|
|
|
|
def __init__(self, beta=0.5, lambda_=0.05):
|
|
super().__init__()
|
|
self.beta = beta
|
|
self.lambda_ = lambda_
|
|
mu = torch.Tensor([-1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75, 1.0])
|
|
self.mu_x = mu.repeat(9)
|
|
self.mu_y = mu.unsqueeze(0).t().repeat(1, 9).view(-1)
|
|
|
|
@staticmethod
|
|
def batch_rSMI(img1, img2, mu_x, mu_y, beta, lambda_):
|
|
assert img1.size() == img2.size()
|
|
|
|
num_pixel = img1.size(1) * img1.size(2) * img2.size(3)
|
|
|
|
mat_k = gaussian_radial_basis_function(img1, mu_x, sigma=1)
|
|
mat_l = gaussian_radial_basis_function(img2, mu_y, sigma=1)
|
|
|
|
mat_k_mul_mat_l = mat_k * mat_l
|
|
h_hat = (1 - beta) * (mat_k_mul_mat_l.matmul(mat_k_mul_mat_l.transpose(1, 2))) / num_pixel
|
|
h_hat += beta * (mat_k.matmul(mat_k.transpose(1, 2)) * mat_l.matmul(mat_l.transpose(1, 2))) / (num_pixel ** 2)
|
|
small_h_hat = mat_k.sum(2, keepdim=True) * mat_l.sum(2, keepdim=True) / (num_pixel ** 2)
|
|
|
|
R = torch.eye(h_hat.size(1)).to(img1.device)
|
|
alpha = (h_hat + lambda_ * R).inverse().matmul(small_h_hat)
|
|
|
|
rSMI = (2 * alpha.transpose(1, 2).matmul(small_h_hat)) - alpha.transpose(1, 2).matmul(h_hat).matmul(alpha) - 1
|
|
return rSMI
|
|
|
|
def forward(self, fake, real):
|
|
rSMI = self.batch_rSMI(fake, real, self.mu_x, self.mu_y, self.beta, self.lambda_)
|
|
return -rSMI.squeeze().mean()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
mg = MGCLoss().to("cuda")
|
|
|
|
|
|
def norm(x):
|
|
x -= x.min()
|
|
x /= x.max()
|
|
return (x - 0.5) * 2
|
|
|
|
|
|
x1 = norm(torch.randn(5, 3, 256, 256))
|
|
x2 = norm(x1 * 2 + 1)
|
|
x3 = norm(torch.randn(5, 3, 256, 256))
|
|
x4 = norm(torch.exp(x3))
|
|
print(mg(x1, x1), mg(x1, x2), mg(x1, x3), mg(x1, x4))
|