File size: 2,285 Bytes
803ef9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import torch
import torch.nn as nn
from torch.nn.functional import conv2d
class Whitening2d(nn.Module):
def __init__(self, num_features, momentum=0.01, track_running_stats=True, eps=0):
super(Whitening2d, self).__init__()
self.num_features = num_features
self.momentum = momentum
self.track_running_stats = track_running_stats
self.eps = eps
if self.track_running_stats:
self.register_buffer(
"running_mean", torch.zeros([1, self.num_features, 1, 1])
)
self.register_buffer("running_variance", torch.eye(self.num_features))
def forward(self, x):
x = x.unsqueeze(2).unsqueeze(3)
m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
if not self.training and self.track_running_stats: # for inference
m = self.running_mean
xn = x - m
T = xn.permute(1, 0, 2, 3).contiguous().view(self.num_features, -1)
f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)
eye = torch.eye(self.num_features).type(f_cov.type())
if not self.training and self.track_running_stats: # for inference
f_cov = self.running_variance
f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye
inv_sqrt = torch.linalg.solve_triangular(
torch.linalg.cholesky(f_cov_shrinked),
eye,
upper=False
)
inv_sqrt = inv_sqrt.contiguous().view(
self.num_features, self.num_features, 1, 1
)
decorrelated = conv2d(xn, inv_sqrt)
if self.training and self.track_running_stats:
self.running_mean = torch.add(
self.momentum * m.detach(),
(1 - self.momentum) * self.running_mean,
out=self.running_mean,
)
self.running_variance = torch.add(
self.momentum * f_cov.detach(),
(1 - self.momentum) * self.running_variance,
out=self.running_variance,
)
return decorrelated.squeeze(2).squeeze(2)
def extra_repr(self):
return "features={}, eps={}, momentum={}".format(
self.num_features, self.eps, self.momentum
)
|