forked from ccf-ai-infra/GPUCodeForces
32 lines
877 B
Python
32 lines
877 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, eps=1e-5):
|
|
super().__init__()
|
|
self.eps = eps
|
|
|
|
def forward(self, img1: torch.Tensor, img2: torch.Tensor, z1: torch.Tensor, z2: torch.Tensor) -> torch.Tensor:
|
|
img_diff = torch.abs(img1 - img2).view(img1.size(0), -1).mean(dim=1)
|
|
z_diff = torch.abs(z1 - z2).view(z1.size(0), -1).mean(dim=1)
|
|
|
|
loss = z_diff / (img_diff + self.eps)
|
|
return loss.mean()
|
|
|
|
|
|
batch_size = 32
|
|
c, h, w = 3, 64, 64
|
|
z_dim = 128
|
|
|
|
|
|
def get_inputs():
|
|
img1 = torch.randn(batch_size, c, h, w, dtype=torch.float32)
|
|
img2 = torch.randn(batch_size, c, h, w, dtype=torch.float32)
|
|
z1 = torch.randn(batch_size, z_dim, dtype=torch.float32)
|
|
z2 = torch.randn(batch_size, z_dim, dtype=torch.float32)
|
|
return [img1, img2, z1, z2]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [1e-5] |