| import torch |
| import torch.nn as nn |
| import torch.nn.functional as functional |
|
|
| try: |
| from queue import Queue |
| except ImportError: |
| from Queue import Queue |
|
|
| from .functions import * |
|
|
|
|
| class ABN(nn.Module): |
| """Activated Batch Normalization |
| |
| This gathers a `BatchNorm2d` and an activation function in a single module |
| """ |
|
|
| def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01): |
| """Creates an Activated Batch Normalization module |
| |
| Parameters |
| ---------- |
| num_features : int |
| Number of feature channels in the input and output. |
| eps : float |
| Small constant to prevent numerical issues. |
| momentum : float |
| Momentum factor applied to compute running statistics as. |
| affine : bool |
| If `True` apply learned scale and shift transformation after normalization. |
| activation : str |
| Name of the activation functions, one of: `leaky_relu`, `elu` or `none`. |
| slope : float |
| Negative slope for the `leaky_relu` activation. |
| """ |
| super(ABN, self).__init__() |
| self.num_features = num_features |
| self.affine = affine |
| self.eps = eps |
| self.momentum = momentum |
| self.activation = activation |
| self.slope = slope |
| if self.affine: |
| self.weight = nn.Parameter(torch.ones(num_features)) |
| self.bias = nn.Parameter(torch.zeros(num_features)) |
| else: |
| self.register_parameter('weight', None) |
| self.register_parameter('bias', None) |
| self.register_buffer('running_mean', torch.zeros(num_features)) |
| self.register_buffer('running_var', torch.ones(num_features)) |
| self.reset_parameters() |
|
|
| def reset_parameters(self): |
| nn.init.constant_(self.running_mean, 0) |
| nn.init.constant_(self.running_var, 1) |
| if self.affine: |
| nn.init.constant_(self.weight, 1) |
| nn.init.constant_(self.bias, 0) |
|
|
| def forward(self, x): |
| x = functional.batch_norm(x, self.running_mean, self.running_var, self.weight, self.bias, |
| self.training, self.momentum, self.eps) |
|
|
| if self.activation == ACT_RELU: |
| return functional.relu(x, inplace=True) |
| elif self.activation == ACT_LEAKY_RELU: |
| return functional.leaky_relu(x, negative_slope=self.slope, inplace=True) |
| elif self.activation == ACT_ELU: |
| return functional.elu(x, inplace=True) |
| else: |
| return x |
|
|
| def __repr__(self): |
| rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \ |
| ' affine={affine}, activation={activation}' |
| if self.activation == "leaky_relu": |
| rep += ', slope={slope})' |
| else: |
| rep += ')' |
| return rep.format(name=self.__class__.__name__, **self.__dict__) |
|
|
|
|
| class InPlaceABN(ABN): |
| """InPlace Activated Batch Normalization""" |
|
|
| def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01): |
| """Creates an InPlace Activated Batch Normalization module |
| |
| Parameters |
| ---------- |
| num_features : int |
| Number of feature channels in the input and output. |
| eps : float |
| Small constant to prevent numerical issues. |
| momentum : float |
| Momentum factor applied to compute running statistics as. |
| affine : bool |
| If `True` apply learned scale and shift transformation after normalization. |
| activation : str |
| Name of the activation functions, one of: `leaky_relu`, `elu` or `none`. |
| slope : float |
| Negative slope for the `leaky_relu` activation. |
| """ |
| super(InPlaceABN, self).__init__(num_features, eps, momentum, affine, activation, slope) |
|
|
| def forward(self, x): |
| x, _, _ = inplace_abn(x, self.weight, self.bias, self.running_mean, self.running_var, |
| self.training, self.momentum, self.eps, self.activation, self.slope) |
| return x |
|
|
|
|
| class InPlaceABNSync(ABN): |
| """InPlace Activated Batch Normalization with cross-GPU synchronization |
| This assumes that it will be replicated across GPUs using the same mechanism as in `nn.DistributedDataParallel`. |
| """ |
|
|
| def forward(self, x): |
| x, _, _ = inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var, |
| self.training, self.momentum, self.eps, self.activation, self.slope) |
| return x |
|
|
| def __repr__(self): |
| rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \ |
| ' affine={affine}, activation={activation}' |
| if self.activation == "leaky_relu": |
| rep += ', slope={slope})' |
| else: |
| rep += ')' |
| return rep.format(name=self.__class__.__name__, **self.__dict__) |
|
|
|
|
|
|
