megengine.module.normalization 源代码

import numpy as np

import megengine as mge
import megengine.functional as F
from megengine import Parameter

from .init import ones_, zeros_
from .module import Module


[文档]class GroupNorm(Module): """Simple implementation of GroupNorm. Only support 4d tensor now. Reference: https://arxiv.org/pdf/1803.08494.pdf. """ def __init__(self, num_groups, num_channels, eps=1e-5, affine=True, **kwargs): super().__init__(**kwargs) assert num_channels % num_groups == 0 self.num_groups = num_groups self.num_channels = num_channels self.eps = eps self.affine = affine if self.affine: self.weight = Parameter(np.ones(num_channels, dtype=np.float32)) self.bias = Parameter(np.zeros(num_channels, dtype=np.float32)) else: self.weight = None self.bias = None self.reset_parameters() def reset_parameters(self): if self.affine: ones_(self.weight) zeros_(self.bias) def forward(self, x): x = F.nn.group_norm( x, self.num_groups, self.affine, self.weight, self.bias, self.eps ) return x def _module_info_string(self) -> str: s = ( "groups={num_groups}, channels={num_channels}, " "eps={eps}, affine={affine}" ) return s.format(**self.__dict__)
[文档]class InstanceNorm(Module): """Simple implementation of InstanceNorm. Only support 4d tensor now. Reference: https://arxiv.org/abs/1607.08022. Note that InstanceNorm equals using GroupNome with num_groups=num_channels. """ def __init__(self, num_channels, eps=1e-05, affine=True, **kwargs): super().__init__(**kwargs) self.num_channels = num_channels self.eps = eps self.affine = affine if self.affine: self.weight = Parameter(np.ones(num_channels, dtype="float32")) self.bias = Parameter(np.zeros(num_channels, dtype="float32")) else: self.weight = None self.bias = None self.reset_parameters() def reset_parameters(self): if self.affine: ones_(self.weight) zeros_(self.bias) def forward(self, x): N, C, H, W = x.shape format = x.format assert C == self.num_channels x = x.reshape(N, C, -1) mean = x.mean(axis=2, keepdims=True) var = (x ** 2).mean(axis=2, keepdims=True) - mean * mean x = (x - mean) / F.sqrt(var + self.eps) x = x.reshape(N, C, H, W) if self.affine: x = self.weight.reshape(1, -1, 1, 1) * x + self.bias.reshape(1, -1, 1, 1) # FIXME(czh): remove this after making it a builtin op. if format == "nhwc": x = mge.amp.convert_tensor_format(x, inplace=False) return x def _module_info_string(self) -> str: s = "channels={num_channels}, eps={eps}, affine={affine}" return s.format(**self.__dict__)
[文档]class LayerNorm(Module): """Simple implementation of LayerNorm. Support tensor of any shape as input. Reference: https://arxiv.org/pdf/1803.08494.pdf. """ def __init__(self, normalized_shape, eps=1e-05, affine=True, **kwargs): super().__init__(**kwargs) if isinstance(normalized_shape, int): normalized_shape = (normalized_shape,) self.normalized_shape = tuple(normalized_shape) self.eps = eps self.affine = affine if self.affine: self.weight = Parameter(np.ones(self.normalized_shape, dtype="float32")) self.bias = Parameter(np.zeros(self.normalized_shape, dtype="float32")) else: self.weight = None self.bias = None self.reset_parameters() def reset_parameters(self): if self.affine: ones_(self.weight) zeros_(self.bias) def forward(self, x): x = F.nn.layer_norm( x, self.normalized_shape, self.affine, self.weight, self.bias, self.eps ) return x def _module_info_string(self) -> str: s = "normalized_shape={normalized_shape}, eps={eps}, affine={affine}" return s.format(**self.__dict__)