forked from ccf-ai-infra/TileOPs-Metax
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import torch
|
|
|
|
from workloads.base import WorkloadBase
|
|
|
|
|
|
def _make_tensors(N, C, spatial, dtype, device="cuda"):
|
|
shape = (N, C, *spatial)
|
|
x = torch.randn(*shape, device=device, dtype=dtype)
|
|
weight = torch.randn(C, device=device, dtype=torch.float32)
|
|
bias = torch.randn(C, device=device, dtype=torch.float32)
|
|
running_mean = torch.zeros(C, device=device, dtype=torch.float32)
|
|
running_var = torch.ones(C, device=device, dtype=torch.float32)
|
|
return x, weight, bias, running_mean, running_var
|
|
|
|
class BatchNormBwdTest(WorkloadBase):
|
|
|
|
def __init__(self, N, C, spatial, dtype):
|
|
self.N = N
|
|
self.C = C
|
|
self.spatial = spatial
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self):
|
|
x, weight, bias, running_mean, running_var = _make_tensors(
|
|
self.N, self.C, self.spatial, self.dtype)
|
|
grad_out = torch.randn_like(x)
|
|
# Need mean/rstd from a forward pass.
|
|
x32 = x.float()
|
|
# Compute mean and rstd via native batch norm internals.
|
|
C = self.C
|
|
L = x32.numel() // C
|
|
x_cl = x32.permute(1, 0, *range(2, x32.ndim)).reshape(C, L).contiguous()
|
|
mean = x_cl.mean(dim=1)
|
|
var = x_cl.var(dim=1, unbiased=False)
|
|
rstd = 1.0 / torch.sqrt(var + 1e-5)
|
|
return grad_out, x, weight, mean, rstd
|
|
|
|
class BatchNormFwdTest(WorkloadBase):
|
|
|
|
def __init__(self, N, C, spatial, dtype, training):
|
|
self.N = N
|
|
self.C = C
|
|
self.spatial = spatial
|
|
self.dtype = dtype
|
|
self.training = training
|
|
|
|
def gen_inputs(self):
|
|
return _make_tensors(self.N, self.C, self.spatial, self.dtype)
|