forked from ccf-ai-infra/GPUCodeForces
57 lines
1003 B
Python
57 lines
1003 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
EPS = 1e-8
|
|
|
|
|
|
class Variance(nn.Module):
|
|
|
|
def __init__(self, unbiased=True):
|
|
super().__init__()
|
|
|
|
self.unbiased = unbiased
|
|
|
|
def forward(self, x):
|
|
|
|
x_flat = x.view(x.size(0), -1)
|
|
|
|
D = x_flat.size(1)
|
|
|
|
x_mean = x_flat.mean(dim=1, keepdim=True)
|
|
|
|
x_centered = x_flat - x_mean
|
|
|
|
ssd_sum = (x_centered ** 2).sum(dim=1)
|
|
|
|
if self.unbiased:
|
|
|
|
divisor = D - 1
|
|
else:
|
|
|
|
divisor = D
|
|
|
|
if divisor <= 0:
|
|
return torch.zeros_like(ssd_sum)
|
|
|
|
variance = ssd_sum / divisor
|
|
|
|
return torch.clamp(variance, min=0)
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, unbiased=True):
|
|
super().__init__()
|
|
self.op = Variance(unbiased=unbiased)
|
|
|
|
def forward(self, x):
|
|
return self.op(x)
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |