forked from ccf-ai-infra/GPUCodeForces
43 lines
1.0 KiB
Python
43 lines
1.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
EPS = 1e-8
|
|
|
|
class PearsonCorrelation(nn.Module):
|
|
def __init__(self, eps=1e-8):
|
|
super().__init__()
|
|
self.eps = eps
|
|
|
|
def forward(self, x, y):
|
|
x_flat = x.view(x.size(0), -1)
|
|
y_flat = y.view(y.size(0), -1)
|
|
|
|
x_mean = x_flat.mean(dim=1, keepdim=True)
|
|
y_mean = y_flat.mean(dim=1, keepdim=True)
|
|
|
|
x_centered = x_flat - x_mean
|
|
y_centered = y_flat - y_mean
|
|
|
|
cov = (x_centered * y_centered).sum(dim=1)
|
|
x_var = (x_centered ** 2).sum(dim=1)
|
|
y_var = (y_centered ** 2).sum(dim=1)
|
|
|
|
denom = torch.sqrt(x_var * y_var)
|
|
return cov / (denom + self.eps)
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.op = PearsonCorrelation(EPS)
|
|
|
|
def forward(self, x, y):
|
|
return self.op(x, y)
|
|
|
|
def get_inputs():
|
|
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
y = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
return [x, y]
|
|
|
|
def get_init_inputs():
|
|
return [] |