# -*- coding: utf-8 -*-
import math
from .module import Module
from ...core import *
from ...functions import batch_norm
from ...autograd import Tensor
[docs]class BatchNorm1d(Module):
'''Applies Batch Normalization over a 2D or 3D input.\n
Args:
num_features (int): C from an expected input of size [N,C] or [N,C,L]
eps (float): a value added to the denominator for numerical stability. Default: 1e-5
momentum (float): the value used for the running_mean and running_std computation. Default:0.1
track_running_stats (bool): if True, this module tracks the running mean and std, else this module always uses batch statistics in both training and eval modes.
Shape:
- Input: [N,C] or [N,C,L]
- Output: same as input
'''
def __init__(self, num_features, eps=1e-5, momentum=0.1, track_running_stats=True):
super().__init__()
self.num_features = num_features
self.eps = eps
self.momentum = momentum
self.track_running_stats = track_running_stats
self.num_params = 2*num_features
self.weight = Tensor(np.random.uniform(0,1,(1,num_features)))
self.bias = Tensor(np.zeros((1,num_features)))
if track_running_stats:
self.mean = Tensor(np.zeros((1,num_features)), requires_grad=False)
self.std = Tensor(np.zeros((1,num_features)), requires_grad=False)
def __repr__(self):
return '{}({}, eps={}, momentum={}, track_running_stats={}) at 0x{:0{}X}'.format(self.__class__.__name__, self.num_features, self.eps, self.momentum, self.track_running_stats, id(self), 16)
[docs] def forward(self, x):
assert x.ndim in [2,3]
if x.ndim == 3 and self.weight.ndim != 3:
self.weight = Tensor(np.random.uniform(0,1,(1,self.num_features,1)))
self.bias = Tensor(np.zeros((1,self.num_features,1)))
if self.track_running_stats:
self.mean = Tensor(np.zeros((1,self.num_features,1)), requires_grad=False)
self.std = Tensor(np.zeros((1,self.num_features,1)), requires_grad=False)
axis = 0 if x.ndim == 2 else (0, 2)
if self.training:
mean = np.mean(x.data, axis=axis, keepdims=True)
std = np.std(x.data, axis=axis, keepdims=True)
if self.track_running_stats:
self.mean.data = self.momentum*mean + (1 - self.momentum)*self.mean.data
self.std.data = self.momentum*std + (1 - self.momentum)*self.std.data
result = batch_norm(x, mean, std, self.weight, self.bias, axis, self.eps)
else:
if self.track_running_stats:
result = batch_norm(x, self.mean.data, self.std.data, self.weight, self.bias, axis, self.eps)
else:
mean = np.mean(x.data, axis=axis, keepdims=True)
std = np.std(x.data, axis=axis, keepdims=True)
result = batch_norm(x, mean, std, self.weight, self.bias, axis, self.eps)
if self.input_shape is None:
self.input_shape = x.shape
if self.output_shape is None:
self.output_shape = result.shape
return result
[docs]class BatchNorm2d(Module):
'''Applies Batch Normalization over a 4D input\n
Args:
num_features (int): C from an expected input of size [N,C,H,W]
eps (float): a value added to the denominator for numerical stability. Default: 1e-5
momentum (float): the value used for the running_mean and running_std computation. Default:0.1
track_running_stats (bool): if True, this module tracks the running mean and std, else this module always uses batch statistics in both training and eval modes.
Shape:
- Input: [N,C,H,W]
- Output: same as input
'''
def __init__(self, num_features, eps=1e-5, momentum=0.1, track_running_stats=True):
super().__init__()
self.num_features = num_features
self.eps = eps
self.momentum = momentum
self.track_running_stats = track_running_stats
self.num_params = 2*num_features
self.weight = Tensor(np.random.uniform(0,1,(1,num_features,1,1)))
self.bias = Tensor(np.zeros((1,num_features,1,1)))
if track_running_stats:
self.mean = Tensor(np.zeros((1,num_features,1,1)), requires_grad=False)
self.std = Tensor(np.zeros((1,num_features,1,1)), requires_grad=False)
def __repr__(self):
return '{}({}, eps={}, momentum={}, track_running_stats={}) at 0x{:0{}X}'.format(self.__class__.__name__, self.num_features, self.eps, self.momentum, self.track_running_stats, id(self), 16)
[docs] def forward(self, x):
axis = (0,2,3)
if self.training:
mean = np.mean(x.data, axis=axis, keepdims=True)
std = np.std(x.data, axis=axis, keepdims=True)
if self.track_running_stats:
self.mean.data = self.momentum*mean + (1 - self.momentum)*self.mean.data
self.std.data = self.momentum*std + (1 - self.momentum)*self.std.data
result = batch_norm(x, mean, std, self.weight, self.bias, axis, self.eps)
else:
if self.track_running_stats:
result = batch_norm(x, self.mean.data, self.std.data, self.weight, self.bias, axis, self.eps)
else:
mean = np.mean(x.data, axis=axis, keepdims=True)
std = np.std(x.data, axis=axis, keepdims=True)
result = batch_norm(x, mean, std, self.weight, self.bias, axis, self.eps)
if self.input_shape is None:
self.input_shape = x.shape
if self.output_shape is None:
self.output_shape = result.shape
return result
[docs]class BatchNorm3d(Module):
'''Applies Batch Normalization over a 5D input\n
Args:
num_features (int): C from an expected input of size [N,C,D,H,W]
eps (float): a value added to the denominator for numerical stability. Default: 1e-5
momentum (float): the value used for the running_mean and running_std computation. Default:0.1
track_running_stats (bool): if True, this module tracks the running mean and std, else this module always uses batch statistics in both training and eval modes.
Shape:
- Input: [N,C,D,H,W]
- Output: same as input
'''
def __init__(self, num_features, eps=1e-5, momentum=0.1, track_running_stats=True):
super().__init__()
self.num_features = num_features
self.eps = eps
self.momentum = momentum
self.track_running_stats = track_running_stats
self.num_params = 2*num_features
self.weight = Tensor(np.random.uniform(0,1,(1,num_features,1,1,1)))
self.bias = Tensor(np.zeros((1,num_features,1,1,1)))
if track_running_stats:
self.mean = Tensor(np.zeros((1,num_features,1,1,1)), requires_grad=False)
self.std = Tensor(np.zeros((1,num_features,1,1,1)), requires_grad=False)
def __repr__(self):
return '{}({}, eps={}, momentum={}, track_running_stats={}) at 0x{:0{}X}'.format(self.__class__.__name__, self.num_features, self.eps, self.momentum, self.track_running_stats, id(self), 16)
[docs] def forward(self, x):
axis = (0,2,3,4)
if self.training:
mean = np.mean(x.data, axis=axis, keepdims=True)
std = np.std(x.data, axis=axis, keepdims=True)
if self.track_running_stats:
self.mean.data = self.momentum*mean + (1 - self.momentum)*self.mean.data
self.std.data = self.momentum*std + (1 - self.momentum)*self.std.data
result = batch_norm(x, mean, std, self.weight, self.bias, axis, self.eps)
else:
if self.track_running_stats:
result = batch_norm(x, self.mean.data, self.std.data, self.weight, self.bias, axis, self.eps)
else:
mean = np.mean(x.data, axis=axis, keepdims=True)
std = np.std(x.data, axis=axis, keepdims=True)
result = batch_norm(x, mean, std, self.weight, self.bias, axis, self.eps)
if self.input_shape is None:
self.input_shape = x.shape
if self.output_shape is None:
self.output_shape = result.shape
return result