mirror of
https://github.com/RVC-Boss/GPT-SoVITS.git
synced 2025-10-07 23:48:48 +08:00
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import logging
|
|
|
|
import torch
|
|
from torch import Tensor, nn
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Normalizer(nn.Module):
|
|
def __init__(self, momentum=0.01, eps=1e-9):
|
|
super().__init__()
|
|
self.momentum = momentum
|
|
self.eps = eps
|
|
self.running_mean_unsafe: Tensor
|
|
self.running_var_unsafe: Tensor
|
|
self.register_buffer("running_mean_unsafe", torch.full([], torch.nan))
|
|
self.register_buffer("running_var_unsafe", torch.full([], torch.nan))
|
|
|
|
@property
|
|
def started(self):
|
|
return not torch.isnan(self.running_mean_unsafe)
|
|
|
|
@property
|
|
def running_mean(self):
|
|
if not self.started:
|
|
return torch.zeros_like(self.running_mean_unsafe)
|
|
return self.running_mean_unsafe
|
|
|
|
@property
|
|
def running_std(self):
|
|
if not self.started:
|
|
return torch.ones_like(self.running_var_unsafe)
|
|
return (self.running_var_unsafe + self.eps).sqrt()
|
|
|
|
@torch.no_grad()
|
|
def _ema(self, a: Tensor, x: Tensor):
|
|
return (1 - self.momentum) * a + self.momentum * x
|
|
|
|
def update_(self, x):
|
|
if not self.started:
|
|
self.running_mean_unsafe = x.mean()
|
|
self.running_var_unsafe = x.var()
|
|
else:
|
|
self.running_mean_unsafe = self._ema(self.running_mean_unsafe, x.mean())
|
|
self.running_var_unsafe = self._ema(self.running_var_unsafe, (x - self.running_mean).pow(2).mean())
|
|
|
|
def forward(self, x: Tensor, update=True):
|
|
if self.training and update:
|
|
self.update_(x)
|
|
self.stats = dict(mean=self.running_mean.item(), std=self.running_std.item())
|
|
x = (x - self.running_mean) / self.running_std
|
|
return x
|
|
|
|
def inverse(self, x: Tensor):
|
|
return x * self.running_std + self.running_mean
|