forked from ccf-ai-infra/GPUCodeForces
88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
# switchablenorm_torch.py
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# 定义输入尺寸 (N, C, H, W)
|
|
N, C, H, W = 16, 64, 32, 32
|
|
EPS = 1e-5
|
|
|
|
class SN(nn.Module):
|
|
|
|
def __init__(self, num_channels, eps):
|
|
super().__init__()
|
|
self.eps = eps
|
|
|
|
self.weight = nn.Parameter(torch.ones(1, num_channels, 1, 1)) # gamma
|
|
self.bias = nn.Parameter(torch.zeros(1, num_channels, 1, 1)) # beta
|
|
|
|
self.w_in = nn.Parameter(torch.ones(num_channels))
|
|
self.w_ln = nn.Parameter(torch.ones(num_channels))
|
|
self.w_bn = nn.Parameter(torch.ones(num_channels))
|
|
|
|
self.register_buffer('running_mean', torch.zeros(num_channels))
|
|
self.register_buffer('running_var', torch.ones(num_channels))
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
ln_mean = x.mean(dim=[1, 2, 3], keepdim=True)
|
|
ln_var = x.var(dim=[1, 2, 3], keepdim=True)
|
|
|
|
|
|
in_mean = x.mean(dim=[2, 3], keepdim=True)
|
|
in_var = x.var(dim=[2, 3], keepdim=True)
|
|
|
|
|
|
bn_mean = self.running_mean.view(1, C, 1, 1)
|
|
bn_var = self.running_var.view(1, C, 1, 1)
|
|
|
|
w_sum = self.w_in.abs() + self.w_ln.abs() + self.w_bn.abs()
|
|
w_in_norm = (self.w_in.abs() / w_sum).view(1, C, 1, 1)
|
|
w_ln_norm = (self.w_ln.abs() / w_sum).view(1, C, 1, 1)
|
|
w_bn_norm = (self.w_bn.abs() / w_sum).view(1, C, 1, 1)
|
|
|
|
mean = w_in_norm * in_mean + w_ln_norm * ln_mean + w_bn_norm * bn_mean
|
|
|
|
|
|
var_in_M2 = in_var + in_mean.pow(2)
|
|
var_ln_M2 = ln_var + ln_mean.pow(2)
|
|
var_bn_M2 = bn_var + bn_mean.pow(2)
|
|
|
|
aggregated_var_M2 = w_in_norm * var_in_M2 + w_ln_norm * var_ln_M2 + w_bn_norm * var_bn_M2
|
|
var = aggregated_var_M2 - mean.pow(2)
|
|
|
|
|
|
x_norm = (x - mean) / torch.sqrt(var + self.eps)
|
|
return x_norm * self.weight + self.bias
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, weight, bias, w_in, w_ln, w_bn):
|
|
super().__init__()
|
|
self.sn = SN(C, EPS)
|
|
|
|
with torch.no_grad():
|
|
self.sn.weight.data.copy_(weight)
|
|
self.sn.bias.data.copy_(bias)
|
|
self.sn.w_in.data.copy_(w_in.squeeze())
|
|
self.sn.w_ln.data.copy_(w_ln.squeeze())
|
|
self.sn.w_bn.data.copy_(w_bn.squeeze())
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.sn(x)
|
|
|
|
|
|
def get_inputs():
|
|
|
|
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
|
|
w_in = torch.ones(C)
|
|
w_ln = torch.ones(C)
|
|
w_bn = torch.ones(C)
|
|
weight = torch.ones(1, C, 1, 1)
|
|
bias = torch.zeros(1, C, 1, 1)
|
|
return [weight, bias, w_in, w_ln, w_bn] |