<commit_before>"""
Batch Normalization: Accelerating Deep Network Training by Reducing Internal
Covariate Shift, http://arxiv.org/abs/1502.03167

This implementation is heavily inspired by code from http://github.com/torch/nn
"""
import numpy as np
import cudarray as ca
from ...base import ParamMixin, PhaseMixin
from ...parameter import Parameter
from ...filler import UniformFiller
from ..base import UnaryElementWise


class BatchNormalization(UnaryElementWise, ParamMixin, PhaseMixin):
    def __init__(self, momentum=0.9, eps=1e-5, affine=True):
        self.momentum = momentum
        self.eps = eps
        self.phase = 'train'
        self.affine = affine
        if self.affine:
            self.gamma = Parameter(UniformFiller(low=0, high=1))
            self.beta = Parameter(0.0)
        self.running_mean = None

    def setup(self):
        super(BatchNormalization, self).setup()
        if len(self.out_shape) != 2:
            raise ValueError('Only 1D data supported')
        reduced_shape = 1, self.out_shape[1]
        if self.running_mean is None:
            self.running_mean = ca.zeros(reduced_shape)
            self.running_std = ca.ones(reduced_shape)
            if self.affine:
                self.gamma.setup(reduced_shape)
                self.beta.setup(reduced_shape)
        else:
            if self.running_mean.shape != reduced_shape:
                raise ValueError('New input shape is not compatible')
        self._tmp_batch_inv_std = ca.zeros(reduced_shape)
        self._tmp_batch_centered = ca.zeros(self.out_shape)

    @property
    def params(self):
        params = []
        if self.affine:
            params = [self.gamma, self.beta]
        return params

    @params.setter
    def params(self, params):
        if self.affine:
            self.gamma, self.beta = params

    def fprop(self):
        if self.phase == 'train':
            # Calculate batch mean
            tmp = ca.mean(self.x.out, axis=0, keepdims=True)
            # Center input
            ca.subtract(self.x.out, tmp, self._tmp_batch_centered)
            # Update running mean
            tmp *= 1 - self.momentum
            self.running_mean *= self.momentum
            self.running_mean += tmp
            # Calculate batch variance
            ca.power(self._tmp_batch_centered, 2, self.out)
            ca.mean(self.out, axis=0, keepdims=True,
                    out=self._tmp_batch_inv_std)
            # Calculate 1 / E([x - E(x)]^2)
            self._tmp_batch_inv_std += self.eps
            ca.sqrt(self._tmp_batch_inv_std, self._tmp_batch_inv_std)
            ca.power(self._tmp_batch_inv_std, -1, self._tmp_batch_inv_std)
            # Normalize input
            ca.multiply(self._tmp_batch_centered, self._tmp_batch_inv_std,
                        self.out)
            # Update running std
            self.running_std *= self.momentum
            ca.multiply(self._tmp_batch_inv_std, 1-self.momentum, tmp)
            self.running_std += tmp
        elif self.phase == 'test':
            ca.subtract(self.x.out, self.running_mean, self.out)
            self.out *= self.running_std
        else:
            raise ValueError('Invalid phase: %s' % self.phase)
        if self.affine:
            self.out *= self.gamma.array
            self.out += self.beta.array

    def bprop(self):
        ca.multiply(self._tmp_batch_centered, self.out_grad, self.x.out_grad)
        tmp = ca.mean(self.x.out_grad, axis=0, keepdims=True)
        ca.multiply(self._tmp_batch_centered, tmp, self.x.out_grad)
        self.x.out_grad *= -1
        self.x.out_grad *= self._tmp_batch_inv_std
        self.x.out_grad *= self._tmp_batch_inv_std

        ca.mean(self.out_grad, axis=0, keepdims=True, out=tmp)
        self.x.out_grad += self.out_grad
        self.x.out_grad -= tmp
        self.x.out_grad *= self._tmp_batch_inv_std

        if self.affine:
            self.x.out_grad *= self.gamma.array
            # Normalized input
            self._tmp_batch_centered *= self._tmp_batch_inv_std
            self._tmp_batch_centered *= self.out_grad
            ca.sum(self._tmp_batch_centered, axis=0, keepdims=True,
                   out=self.gamma.grad_array)
            ca.sum(self.out_grad, axis=0, keepdims=True,
                   out=self.beta.grad_array)


class SpatialBatchNormalization(UnaryElementWise, ParamMixin, PhaseMixin):
    def __init__(self, momentum=0.9, eps=1e-5, affine=True):
        self.momentum = momentum
        self.eps = eps
        self.phase = 'train'
        self.affine = affine
        if self.affine:
            self.gamma = Parameter(UniformFiller(low=0, high=1))
            self.beta = Parameter(0.0)
        self.running_mean = None

    def setup(self):
        super(SpatialBatchNormalization, self).setup()
        if len(self.out_shape) != 4:
            raise ValueError('Only 4D data supported')
        reduced_shape = 1, self.out_shape[1], 1, 1
        if self.running_mean is None:
            self.running_mean = ca.zeros(reduced_shape)
            self.running_std = ca.ones(reduced_shape)
            if self.affine:
                self.gamma.setup(reduced_shape)
                self.beta.setup(reduced_shape)
        else:
            if self.running_mean.shape != reduced_shape:
                raise ValueError('New input shape is not compatible')
        self._tmp_batch_inv_std = ca.zeros(reduced_shape)
        self._tmp_batch_centered = ca.zeros(self.out_shape)

    @property
    def params(self):
        params = []
        if self.affine:
            params = [self.gamma, self.beta]
        return params

    @params.setter
    def params(self, params):
        if self.affine:
            self.gamma, self.beta = params

    def fprop(self):
        if self.phase == 'train':
            # Calculate batch mean
            tmp = ca.mean(ca.mean(self.x.out, axis=0, keepdims=True),
                          axis=(2, 3), keepdims=True)
            # Center input
            ca.subtract(self.x.out, tmp, self._tmp_batch_centered)
            # Update running mean
            tmp *= 1 - self.momentum
            self.running_mean *= self.momentum
            self.running_mean += tmp
            # Calculate batch variance
            ca.power(self._tmp_batch_centered, 2, self.out)
            ca.mean(ca.mean(self.out, axis=0, keepdims=True), axis=(2, 3),
                    keepdims=True, out=self._tmp_batch_inv_std)
            # Calculate 1 / E([x - E(x)]^2)
            self._tmp_batch_inv_std += self.eps
            ca.sqrt(self._tmp_batch_inv_std, self._tmp_batch_inv_std)
            ca.power(self._tmp_batch_inv_std, -1, self._tmp_batch_inv_std)
            # Normalize input
            ca.multiply(self._tmp_batch_centered, self._tmp_batch_inv_std,
                        self.out)
            # Update running std
            self.running_std *= self.momentum
            ca.multiply(self._tmp_batch_inv_std, 1-self.momentum, tmp)
            self.running_std += tmp
        elif self.phase == 'test':
            ca.subtract(self.x.out, self.running_mean, self.out)
            self.out *= self.running_std
        else:
            raise ValueError('Invalid phase: %s' % self.phase)
        if self.affine:
            self.out *= self.gamma.array
            self.out += self.beta.array

    def bprop(self):
        ca.multiply(self._tmp_batch_centered, self.out_grad, self.x.out_grad)
        tmp = ca.mean(ca.mean(self.x.out_grad, axis=0, keepdims=True),
                      axis=(2, 3), keepdims=True)
        ca.multiply(self._tmp_batch_centered, tmp, self.x.out_grad)
        self.x.out_grad *= -1
        self.x.out_grad *= self._tmp_batch_inv_std
        self.x.out_grad *= self._tmp_batch_inv_std

        tmp = ca.mean(ca.mean(self.out_grad, axis=0, keepdims=True),
                      axis=(2, 3), keepdims=True)
        self.x.out_grad += self.out_grad
        self.x.out_grad -= tmp
        self.x.out_grad *= self._tmp_batch_inv_std

        if self.affine:
            self.x.out_grad *= self.gamma.array
            # Normalized input
            self._tmp_batch_centered *= self._tmp_batch_inv_std
            self._tmp_batch_centered *= self.out_grad<commit_msg>add noise_std parameter<commit_after>"""
Batch Normalization: Accelerating Deep Network Training by Reducing Internal
Covariate Shift, http://arxiv.org/abs/1502.03167

This implementation is heavily inspired by code from http://github.com/torch/nn
"""
import cudarray as ca
from ...base import ParamMixin, PhaseMixin
from ...parameter import Parameter
from ...filler import UniformFiller
from ..base import UnaryElementWise


class BatchNormalization(UnaryElementWise, ParamMixin, PhaseMixin):
    def __init__(self, momentum=0.9, eps=1e-5, noise_std=0.0, affine=True):
        self.momentum = momentum
        self.eps = eps
        self.phase = 'train'
        self.affine = affine
        if self.affine:
            self.gamma = Parameter(UniformFiller(low=0, high=1))
            self.beta = Parameter(0.0)
        self.running_mean = None
        self.noise_std = noise_std

    def setup(self):
        super(BatchNormalization, self).setup()
        reduced_shape = (1,) + self.out_shape[1:]
        if self.running_mean is None:
            self.running_mean = ca.zeros(reduced_shape)
            self.running_std = ca.ones(reduced_shape)
            if self.affine:
                self.gamma.setup(reduced_shape)
                self.beta.setup(reduced_shape)
        else:
            if self.running_mean.shape != reduced_shape:
                raise ValueError('New input shape is not compatible')
        self._tmp_batch_inv_std = ca.zeros(reduced_shape)
        self._tmp_batch_centered = ca.zeros(self.out_shape)

    @property
    def params(self):
        params = []
        if self.affine:
            params = [self.gamma, self.beta]
        return params

    @params.setter
    def params(self, params):
        if self.affine:
            self.gamma, self.beta = params

    def fprop(self):
        if self.phase == 'train':
            # Calculate batch mean
            tmp = ca.mean(self.x.out, axis=0, keepdims=True)
            # Center input
            ca.subtract(self.x.out, tmp, self._tmp_batch_centered)
            # Update running mean
            tmp *= 1 - self.momentum
            self.running_mean *= self.momentum
            self.running_mean += tmp
            # Calculate batch variance
            ca.power(self._tmp_batch_centered, 2, self.out)
            ca.mean(self.out, axis=0, keepdims=True,
                    out=self._tmp_batch_inv_std)
            # Calculate 1 / E([x - E(x)]^2)
            self._tmp_batch_inv_std += self.eps
            ca.sqrt(self._tmp_batch_inv_std, self._tmp_batch_inv_std)
            ca.power(self._tmp_batch_inv_std, -1, self._tmp_batch_inv_std)
            # Normalize input
            ca.multiply(self._tmp_batch_centered, self._tmp_batch_inv_std,
                        self.out)
            # Update running std
            self.running_std *= self.momentum
            ca.multiply(self._tmp_batch_inv_std, 1-self.momentum, tmp)
            self.running_std += tmp

            if self.noise_std > 0.0:
                noise = ca.random.normal(scale=self.noise_std,
                                         size=self.out_shape)
                ca.add(self.out, noise, self.out)
        elif self.phase == 'test':
            ca.subtract(self.x.out, self.running_mean, self.out)
            self.out *= self.running_std
        else:
            raise ValueError('Invalid phase: %s' % self.phase)
        if self.affine:
            self.out *= self.gamma.array
            self.out += self.beta.array

    def bprop(self):
        ca.multiply(self._tmp_batch_centered, self.out_grad, self.x.out_grad)
        tmp = ca.mean(self.x.out_grad, axis=0, keepdims=True)
        ca.multiply(self._tmp_batch_centered, tmp, self.x.out_grad)
        self.x.out_grad *= -1
        self.x.out_grad *= self._tmp_batch_inv_std
        self.x.out_grad *= self._tmp_batch_inv_std

        ca.mean(self.out_grad, axis=0, keepdims=True, out=tmp)
        self.x.out_grad += self.out_grad
        self.x.out_grad -= tmp
        self.x.out_grad *= self._tmp_batch_inv_std

        if self.affine:
            self.x.out_grad *= self.gamma.array
            # Normalized input
            self._tmp_batch_centered *= self._tmp_batch_inv_std
            self._tmp_batch_centered *= self.out_grad
            ca.sum(self._tmp_batch_centered, axis=0, keepdims=True,
                   out=self.gamma.grad_array)
            ca.sum(self.out_grad, axis=0, keepdims=True,
                   out=self.beta.grad_array)


class SpatialBatchNormalization(UnaryElementWise, ParamMixin, PhaseMixin):
    def __init__(self, momentum=0.9, eps=1e-5, noise_std=0.0, affine=True):
        self.momentum = momentum
        self.eps = eps
        self.phase = 'train'
        self.affine = affine
        if self.affine:
            self.gamma = Parameter(UniformFiller(low=0, high=1))
            self.beta = Parameter(0.0)
        self.running_mean = None
        self.noise_std = noise_std

    def setup(self):
        super(SpatialBatchNormalization, self).setup()
        if len(self.out_shape) != 4:
            raise ValueError('Only 4D data supported')
        reduced_shape = 1, self.out_shape[1], 1, 1
        if self.running_mean is None:
            self.running_mean = ca.zeros(reduced_shape)
            self.running_std = ca.ones(reduced_shape)
            if self.affine:
                self.gamma.setup(reduced_shape)
                self.beta.setup(reduced_shape)
        else:
            if self.running_mean.shape != reduced_shape:
                raise ValueError('New input shape is not compatible')
        self._tmp_batch_inv_std = ca.zeros(reduced_shape)
        self._tmp_batch_centered = ca.zeros(self.out_shape)

    @property
    def params(self):
        params = []
        if self.affine:
            params = [self.gamma, self.beta]
        return params

    @params.setter
    def params(self, params):
        if self.affine:
            self.gamma, self.beta = params

    def fprop(self):
        if self.phase == 'train':
            # Calculate batch mean
            tmp = ca.mean(ca.mean(self.x.out, axis=0, keepdims=True),
                          axis=(2, 3), keepdims=True)
            # Center input
            ca.subtract(self.x.out, tmp, self._tmp_batch_centered)
            # Update running mean
            tmp *= 1 - self.momentum
            self.running_mean *= self.momentum
            self.running_mean += tmp
            # Calculate batch variance
            ca.power(self._tmp_batch_centered, 2, self.out)
            ca.mean(ca.mean(self.out, axis=0, keepdims=True), axis=(2, 3),
                    keepdims=True, out=self._tmp_batch_inv_std)
            # Calculate 1 / E([x - E(x)]^2)
            self._tmp_batch_inv_std += self.eps
            ca.sqrt(self._tmp_batch_inv_std, self._tmp_batch_inv_std)
            ca.power(self._tmp_batch_inv_std, -1, self._tmp_batch_inv_std)
            # Normalize input
            ca.multiply(self._tmp_batch_centered, self._tmp_batch_inv_std,
                        self.out)
            # Update running std
            self.running_std *= self.momentum
            ca.multiply(self._tmp_batch_inv_std, 1-self.momentum, tmp)
            self.running_std += tmp

            if self.noise_std > 0.0:
                noise = ca.random.normal(scale=self.noise_std,
                                         size=self.out_shape)
                ca.add(self.out, noise, self.out)

        elif self.phase == 'test':
            ca.subtract(self.x.out, self.running_mean, self.out)
            self.out *= self.running_std
        else:
            raise ValueError('Invalid phase: %s' % self.phase)
        if self.affine:
            self.out *= self.gamma.array
            self.out += self.beta.array

    def bprop(self):
        ca.multiply(self._tmp_batch_centered, self.out_grad, self.x.out_grad)
        tmp = ca.mean(ca.mean(self.x.out_grad, axis=0, keepdims=True),
                      axis=(2, 3), keepdims=True)
        ca.multiply(self._tmp_batch_centered, tmp, self.x.out_grad)
        self.x.out_grad *= -1
        self.x.out_grad *= self._tmp_batch_inv_std
        self.x.out_grad *= self._tmp_batch_inv_std

        tmp = ca.mean(ca.mean(self.out_grad, axis=0, keepdims=True),
                      axis=(2, 3), keepdims=True)
        self.x.out_grad += self.out_grad
        self.x.out_grad -= tmp
        self.x.out_grad *= self._tmp_batch_inv_std

        if self.affine:
            self.x.out_grad *= self.gamma.array
            # Normalized input
            self._tmp_batch_centered *= self._tmp_batch_inv_std
            self._tmp_batch_centered *= self.out_grad
{'fraction_non_alphanumeric': 0.10414566158590831, 'fraction_numerical': 0.008303184864480161, 'mean_word_length': 1.9577267146114716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5655166', 'n_tokens_mistral': 5353, 'n_tokens_neox': 5167, 'n_words': 987}