forked from ccf-ai-infra/GPUCodeForces
48 lines
1.0 KiB
Python
48 lines
1.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
EPS = 1e-6
|
|
|
|
|
|
class MinkowskiDistance(nn.Module):
|
|
|
|
def __init__(self, p=2.0, keepdim=False, eps=1e-6):
|
|
super().__init__()
|
|
self.p = p
|
|
self.keepdim = keepdim
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
diff = torch.abs(x - y)
|
|
|
|
pow_diff = torch.pow(diff, self.p)
|
|
|
|
sum_pow = torch.sum(pow_diff, dim=[1, 2, 3], keepdim=self.keepdim)
|
|
|
|
output = torch.pow(sum_pow + self.eps, 1.0 / self.p)
|
|
|
|
return output
|
|
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self, p=3.0):
|
|
super().__init__()
|
|
self.op = MinkowskiDistance(p=p, keepdim=False, eps=EPS)
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
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():
|
|
p = 3.0
|
|
return [p]
|