Live2d Init

live2d/tha3/nn/normalization.py

@@ -0,0 +1,126 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+from torch import layer_norm
+from torch.nn import Module, BatchNorm2d, InstanceNorm2d, Parameter
+from torch.nn.init import normal_, constant_
+
+from tha3.nn.pass_through import PassThrough
+
+
+class PixelNormalization(Module):
+    def __init__(self, epsilon=1e-8):
+        super().__init__()
+        self.epsilon = epsilon
+
+    def forward(self, x):
+        return x / torch.sqrt((x ** 2).mean(dim=1, keepdim=True) + self.epsilon)
+
+
+class NormalizationLayerFactory(ABC):
+    def __init__(self):
+        super().__init__()
+
+    @abstractmethod
+    def create(self, num_features: int, affine: bool = True) -> Module:
+        pass
+
+    @staticmethod
+    def resolve_2d(factory: Optional['NormalizationLayerFactory']) -> 'NormalizationLayerFactory':
+        if factory is None:
+            return InstanceNorm2dFactory()
+        else:
+            return factory
+
+
+class Bias2d(Module):
+    def __init__(self, num_features: int):
+        super().__init__()
+        self.num_features = num_features
+        self.bias = Parameter(torch.zeros(1, num_features, 1, 1))
+
+    def forward(self, x):
+        return x + self.bias
+
+
+class NoNorm2dFactory(NormalizationLayerFactory):
+    def __init__(self):
+        super().__init__()
+
+    def create(self, num_features: int, affine: bool = True) -> Module:
+        if affine:
+            return Bias2d(num_features)
+        else:
+            return PassThrough()
+
+
+class BatchNorm2dFactory(NormalizationLayerFactory):
+    def __init__(self,
+                 weight_mean: Optional[float] = None,
+                 weight_std: Optional[float] = None,
+                 bias: Optional[float] = None):
+        super().__init__()
+        self.bias = bias
+        self.weight_std = weight_std
+        self.weight_mean = weight_mean
+
+    def get_weight_mean(self):
+        if self.weight_mean is None:
+            return 1.0
+        else:
+            return self.weight_mean
+
+    def get_weight_std(self):
+        if self.weight_std is None:
+            return 0.02
+        else:
+            return self.weight_std
+
+    def create(self, num_features: int, affine: bool = True) -> Module:
+        module = BatchNorm2d(num_features=num_features, affine=affine)
+        if affine:
+            if self.weight_mean is not None or self.weight_std is not None:
+                normal_(module.weight, self.get_weight_mean(), self.get_weight_std())
+            if self.bias is not None:
+                constant_(module.bias, self.bias)
+        return module
+
+
+class InstanceNorm2dFactory(NormalizationLayerFactory):
+    def __init__(self):
+        super().__init__()
+
+    def create(self, num_features: int, affine: bool = True) -> Module:
+        return InstanceNorm2d(num_features=num_features, affine=affine)
+
+
+class PixelNormFactory(NormalizationLayerFactory):
+    def __init__(self):
+        super().__init__()
+
+    def create(self, num_features: int, affine: bool = True) -> Module:
+        return PixelNormalization()
+
+
+class LayerNorm2d(Module):
+    def __init__(self, channels: int, affine: bool = True):
+        super(LayerNorm2d, self).__init__()
+        self.channels = channels
+        self.affine = affine
+
+        if self.affine:
+            self.weight = Parameter(torch.ones(1, channels, 1, 1))
+            self.bias = Parameter(torch.zeros(1, channels, 1, 1))
+
+    def forward(self, x):
+        shape = x.size()[1:]
+        y = layer_norm(x, shape) * self.weight + self.bias
+        return y
+
+class LayerNorm2dFactory(NormalizationLayerFactory):
+    def __init__(self):
+        super().__init__()
+
+    def create(self, num_features: int, affine: bool = True) -> Module:
+        return LayerNorm2d(channels=num_features, affine=affine)