# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey) # # See ../../../../LICENSE for clarification regarding multiple authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import collections from itertools import repeat from typing import Optional, Tuple, Union from functools import reduce import logging from torch.cuda.amp import custom_fwd, custom_bwd import random import torch import math import torch.nn as nn import torch.nn.functional as F from torch import Tensor from torch.nn import Embedding as ScaledEmbedding class ScheduledFloat(torch.nn.Module): """ This object is a torch.nn.Module only because we want it to show up in [top_level module].modules(); it does not have a working forward() function. You are supposed to cast it to float, as in, float(parent_module.whatever), and use it as something like a dropout prob. It is a floating point value whose value changes depending on the batch count of the training loop. It is a piecewise linear function where you specifiy the (x,y) pairs in sorted order on x; x corresponds to the batch index. For batch-index values before the first x or after the last x, we just use the first or last y value. Example: self.dropout = ScheduledFloat((0.0, 0.2), (4000.0, 0.0), default=0.0) `default` is used when self.batch_count is not set or in training or mode or in torch.jit scripting mode. """ def __init__(self, *args, default: float = 0.0): super().__init__() # self.batch_count and self.name will be written to in the training loop. self.batch_count = None self.name = None self.default = default assert len(args) >= 1 for (x,y) in args: assert x >= 0 for i in range(len(args) - 1): assert args[i + 1] > args[i], args self.schedule = args def extra_repr(self) -> str: return 'batch_count={}, schedule={}'.format(self.batch_count, self.schedule) def __float__(self): print_prob = 0.0002 def maybe_print(ans): if random.random() < print_prob: logging.info(f"ScheduledFloat: name={self.name}, batch_count={self.batch_count}, ans={ans}") batch_count = self.batch_count if batch_count is None or not self.training or torch.jit.is_scripting(): return float(self.default) if batch_count <= self.schedule[0][0]: ans = self.schedule[0][1] maybe_print(ans) return float(ans) elif batch_count >= self.schedule[-1][0]: ans = self.schedule[-1][1] maybe_print(ans) return float(ans) else: cur_x, cur_y = self.schedule[0] for i in range(1, len(self.schedule)): next_x, next_y = self.schedule[i] if batch_count >= cur_x and batch_count <= next_x: ans = cur_y + (next_y - cur_y) * (batch_count - cur_x) / (next_x - cur_x) maybe_print(ans) return float(ans) cur_x, cur_y = next_x, next_y assert False FloatLike = Union[float, ScheduledFloat] class ActivationBalancerFunction(torch.autograd.Function): @staticmethod def forward( ctx, x: Tensor, scale_factor: Tensor, mean: Tensor, sign_factor: Optional[Tensor], channel_dim: int, ) -> Tensor: if channel_dim < 0: channel_dim += x.ndim ctx.channel_dim = channel_dim xgtmean = (x > mean) if sign_factor is None: ctx.save_for_backward(xgtmean, scale_factor) else: ctx.save_for_backward(xgtmean, scale_factor, sign_factor) return x @staticmethod def backward( ctx, x_grad: Tensor ) -> Tuple[Tensor, None, None, None]: if len(ctx.saved_tensors) == 3: xgtmean, scale_factor, sign_factor = ctx.saved_tensors for _ in range(ctx.channel_dim, x_grad.ndim - 1): scale_factor = scale_factor.unsqueeze(-1) sign_factor = sign_factor.unsqueeze(-1) factor = sign_factor + scale_factor * (xgtmean.to(x_grad.dtype) - 0.5) else: xgtmean, scale_factor = ctx.saved_tensors for _ in range(ctx.channel_dim, x_grad.ndim - 1): scale_factor = scale_factor.unsqueeze(-1) factor = scale_factor * (xgtmean.to(x_grad.dtype) - 0.5) neg_delta_grad = x_grad.abs() * factor return x_grad - neg_delta_grad, None, None, None, None def _compute_scale_factor(x: Tensor, channel_dim: int, min_abs: float, max_abs: float, gain_factor: float, max_factor: float) -> Tuple[Tensor, Tensor]: """ Computes a factor used in ActivationBalancer, that dictates how much we penalize (or anti-penalize) the scale on the features. Returns: (scale_factor, mean) dim. scale_factor: can be positive or negative, between -max_factor and max_factor; dictates penalty or anti-penalty. It is of shape (num_channels,) mean: mean per channel that we use for purposes of scale_factor; actually is clamped to -min_abs..min_abs. Its like (1, num_channels, 1, 1) depending on the shape of x and channel-dim. """ if channel_dim < 0: channel_dim += x.ndim sum_dims = [d for d in range(x.ndim) if d != channel_dim] x_mean = torch.mean(x, dim=sum_dims, keepdim=True).to(torch.float32) # the idea is that for purposes of applying max_abs, we regress effectively # toward zero (assuming min_abs is much less than max_abs). x_mean = x_mean.clamp(min=-min_abs, max=min_abs) x_abs_mean = torch.mean((x - x_mean).abs(), dim=sum_dims).to(torch.float32) if min_abs == 0.0: below_threshold = 0.0 else: # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if # x_abs)_mean , min_abs. below_threshold = ((min_abs - x_abs_mean) * (gain_factor / min_abs)).clamp(min=0, max=max_factor) above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(min=0, max=max_factor) return below_threshold - above_threshold, x_mean def _compute_sign_factor(x: Tensor, channel_dim: int, min_positive: float, max_positive: float, gain_factor: float, max_factor: float) -> Tensor: if channel_dim < 0: channel_dim += x.ndim sum_dims = [d for d in range(x.ndim) if d != channel_dim] proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims) if min_positive == 0.0: factor1 = 0.0 else: # 0 if proportion_positive >= min_positive, else can be # as large as max_factor. factor1 = ((min_positive - proportion_positive) * (gain_factor / min_positive)).clamp_(min=0, max=max_factor) if max_positive == 1.0: factor2 = 0.0 else: # 0 if self.proportion_positive <= max_positive, else can be # as large as -max_factor. factor2 = ((proportion_positive - max_positive) * (gain_factor / (1.0 - max_positive))).clamp_(min=0, max=max_factor) sign_factor = factor1 - factor2 # require min_positive != 0 or max_positive != 1: assert not isinstance(sign_factor, float) return sign_factor def random_cast_to_half(x: Tensor, min_abs: float = 5.0e-06) -> Tensor: """ A randomized way of casting a floating point value to half precision. """ if x.dtype == torch.float16: return x x_abs = x.abs() is_too_small = (x_abs < min_abs) # for elements where is_too_small is true, random_val will contain +-min_abs with # probability (x.abs() / min_abs), and 0.0 otherwise. [so this preserves expectations, # for those elements]. random_val = min_abs * x.sign() * (torch.rand_like(x) * min_abs < x_abs) return torch.where(is_too_small, random_val, x).to(torch.float16) class CutoffEstimator: """ Estimates cutoffs of an arbitrary numerical quantity such that a specified proportion of items will be above the cutoff on average. p is the proportion of items that should be above the cutoff. """ def __init__(self, p: float): self.p = p # total count of items self.count = 0 # total count of items that were above the cutoff self.count_above = 0 # initial cutoff value self.cutoff = 0 def __call__(self, x: float) -> bool: """ Returns true if x is above the cutoff. """ ans = (x > self.cutoff) self.count += 1 if ans: self.count_above += 1 cur_p = self.count_above / self.count delta_p = cur_p - self.p if (delta_p > 0) == ans: q = abs(delta_p) self.cutoff = x * q + self.cutoff * (1-q) return ans class CachingEvalFunction(torch.autograd.Function): # @custom_fwd and @custom_bwd related to automatic mixed precision (amp) an ensure # that the backward path runs with the same autocast context as the forward pass. @staticmethod @custom_fwd def forward(ctx, x: Tensor, m) -> Tensor: """ m might be an nn.Module """ ctx.x_requires_grad = x.requires_grad ctx.m = m # we need any random numbers used in this evaluation and the next evaluation to be identical. # Caution: this assumes you are not going to use any random numbers from torch (for any purpose # that matters in the forward pass), e.g. there should be no dropout. ctx.random_state = random.getstate() # we are inside torch.no_grad() here, so the following won't create the computation graph. with torch.no_grad(): y = m(x) ctx.save_for_backward(x, y) return y @staticmethod @custom_bwd def backward(ctx, y_grad: Tensor): x, y = ctx.saved_tensors x = x.detach() x.requires_grad = ctx.x_requires_grad m = ctx.m # e.g. a nn.Module random_state = random.getstate() # set the state to what we used in the 1st forward pass. random.setstate(ctx.random_state) with torch.enable_grad(): y2 = m(x) assert torch.allclose(y, y2, atol=1.0e-02) # this call to backward() should create grads in the module's parameters y2.backward(gradient=y_grad) # restore the state from before we entered this function random.setstate(random_state) return x.grad, None # x.grad will be None if x.requires_grad is False. def caching_eval(x: Tensor, m: nn.Module) -> Tensor: if m.training: # The purpose of this code is to make all parameters of m reachable in # the computation graph, so that if we give find_unused_parameters=True # to PyTorch's autograd code it does not assign them zero gradient. tot = 0.0 for p in m.parameters(): tot = tot + 0.0 * p.flatten()[0] x = x + tot # tot will be 0, this does nothing. return CachingEvalFunction.apply(x, m) class RandomGradFunction(torch.autograd.Function): """ Does nothing in forward pass; in backward pass, gets rid of very small grads using randomized approach that preserves expectations (intended to reduce roundoff). """ @staticmethod def forward(ctx, x: Tensor, min_abs: float) -> Tensor: ctx.min_abs = min_abs return x @staticmethod def backward(ctx, ans_grad: Tensor) -> Tuple[Tensor, None]: if ans_grad.dtype == torch.float16: return random_cast_to_half(ans_grad.to(torch.float32), min_abs=ctx.min_abs), None else: return ans_grad, None class RandomGrad(torch.nn.Module): """ Gets rid of very small gradients using an expectation-preserving method, intended to increase accuracy of training when using amp (automatic mixed precision) """ def __init__(self, min_abs: float = 5.0e-06): super(RandomGrad, self).__init__() self.min_abs = min_abs def forward(self, x: Tensor): if torch.jit.is_scripting() or not self.training: return x else: return RandomGradFunction.apply(x, self.min_abs) class SoftmaxFunction(torch.autograd.Function): """ Tries to handle half-precision derivatives in a randomized way that should be more accurate for training than the default behavior. """ @staticmethod def forward(ctx, x: Tensor, dim: int): ans = x.softmax(dim=dim) # if x dtype is float16, x.softmax() returns a float32 because # (presumably) that op does not support float16, and autocast # is enabled. if torch.is_autocast_enabled(): ans = ans.to(torch.float16) ctx.save_for_backward(ans) ctx.x_dtype = x.dtype ctx.dim = dim return ans @staticmethod def backward(ctx, ans_grad: Tensor): ans, = ctx.saved_tensors with torch.cuda.amp.autocast(enabled=False): ans_grad = ans_grad.to(torch.float32) ans = ans.to(torch.float32) x_grad = ans_grad * ans x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True) return x_grad, None def softmax(x: Tensor, dim: int): return SoftmaxFunction.apply(x, dim) class MaxEigLimiterFunction(torch.autograd.Function): @staticmethod def forward( ctx, x: Tensor, coeffs: Tensor, direction: Tensor, channel_dim: int, grad_scale: float) -> Tensor: ctx.channel_dim = channel_dim ctx.grad_scale = grad_scale ctx.save_for_backward(x.detach(), coeffs.detach(), direction.detach()) return x @staticmethod def backward(ctx, x_grad, *args): with torch.enable_grad(): (x_orig, coeffs, new_direction) = ctx.saved_tensors x_orig.requires_grad = True num_channels = x_orig.shape[ctx.channel_dim] x = x_orig.transpose(ctx.channel_dim, -1).reshape(-1, num_channels) new_direction.requires_grad = False x = x - x.mean(dim=0) x_var = (x ** 2).mean() x_residual = x - coeffs * new_direction x_residual_var = (x_residual ** 2).mean() # `variance_proportion` is the proportion of the variance accounted for # by the top eigen-direction. This is to be minimized. variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20) variance_proportion.backward() x_orig_grad = x_orig.grad x_extra_grad = x_orig.grad * ctx.grad_scale * x_grad.norm() / (x_orig_grad.norm() + 1.0e-20) return x_grad + x_extra_grad.detach(), None, None, None, None class BasicNormFunction(torch.autograd.Function): # This computes: # scales = (torch.mean((x - bias) ** 2, keepdim=True)) ** -0.5 * log_scale.exp() # return (x - bias) * scales # (after unsqueezing the bias), but it does it in a memory-efficient way so that # it can just store the returned value (chances are, this will also be needed for # some other reason, related to the next operation, so we can save memory). @staticmethod def forward(ctx, x: Tensor, bias: Tensor, log_scale: Tensor, channel_dim: int, store_output_for_backprop: bool) -> Tensor: assert bias.ndim == 1 if channel_dim < 0: channel_dim = channel_dim + x.ndim ctx.store_output_for_backprop = store_output_for_backprop ctx.channel_dim = channel_dim for _ in range(channel_dim + 1, x.ndim): bias = bias.unsqueeze(-1) scales = (torch.mean((x - bias) ** 2, dim=channel_dim, keepdim=True) ** -0.5) * log_scale.exp() ans = x * scales ctx.save_for_backward(ans.detach() if store_output_for_backprop else x, scales.detach(), bias.detach(), log_scale.detach()) return ans @staticmethod def backward(ctx, ans_grad: Tensor) -> Tensor: ans_or_x, scales, bias, log_scale = ctx.saved_tensors if ctx.store_output_for_backprop: x = ans_or_x / scales else: x = ans_or_x x = x.detach() x.requires_grad = True bias.requires_grad = True log_scale.requires_grad = True with torch.enable_grad(): # recompute scales from x, bias and log_scale. scales = (torch.mean((x - bias) ** 2, dim=ctx.channel_dim, keepdim=True) ** -0.5) * log_scale.exp() ans = x * scales ans.backward(gradient=ans_grad) return x.grad, bias.grad.flatten(), log_scale.grad, None, None class BasicNorm(torch.nn.Module): """ This is intended to be a simpler, and hopefully cheaper, replacement for LayerNorm. The observation this is based on, is that Transformer-type networks, especially with pre-norm, sometimes seem to set one of the feature dimensions to a large constant value (e.g. 50), which "defeats" the LayerNorm because the output magnitude is then not strongly dependent on the other (useful) features. Presumably the weight and bias of the LayerNorm are required to allow it to do this. So the idea is to introduce this large constant value as an explicit parameter, that takes the role of the "eps" in LayerNorm, so the network doesn't have to do this trick. We make the "eps" learnable. Args: num_channels: the number of channels, e.g. 512. channel_dim: the axis/dimension corresponding to the channel, interprted as an offset from the input's ndim if negative. shis is NOT the num_channels; it should typically be one of {-2, -1, 0, 1, 2, 3}. log_scale: the initial log-scale that we multiply the output by; this is learnable. log_scale_min: FloatLike, minimum allowed value of log_scale log_scale_max: FloatLike, maximum allowed value of log_scale store_output_for_backprop: only possibly affects memory use; recommend to set to True if you think the output of this module is more likely than the input of this module to be required to be stored for the backprop. """ def __init__( self, num_channels: int, channel_dim: int = -1, # CAUTION: see documentation. log_scale: float = 1.0, log_scale_min: float = -1.5, log_scale_max: float = 1.5, store_output_for_backprop: bool = False ) -> None: super(BasicNorm, self).__init__() self.num_channels = num_channels self.channel_dim = channel_dim self.log_scale = nn.Parameter(torch.tensor(log_scale)) self.bias = nn.Parameter(torch.zeros(num_channels)) self.log_scale_min = log_scale_min self.log_scale_max = log_scale_max self.store_output_for_backprop = store_output_for_backprop def forward(self, x: Tensor) -> Tensor: assert x.shape[self.channel_dim] == self.num_channels if torch.jit.is_scripting(): channel_dim = self.channel_dim if channel_dim < 0: channel_dim += x.ndim bias = self.bias for _ in range(channel_dim + 1, x.ndim): bias = bias.unsqueeze(-1) scales = ((torch.mean((x - bias) ** 2, dim=channel_dim, keepdim=True) ** -0.5) * self.log_scale.exp()) return x * scales log_scale = limit_param_value(self.log_scale, min=float(self.log_scale_min), max=float(self.log_scale_max), training=self.training) return BasicNormFunction.apply(x, self.bias, log_scale, self.channel_dim, self.store_output_for_backprop) class PositiveConv1d(nn.Conv1d): """ A modified form of nn.Conv1d where the weight parameters are constrained to be positive and there is no bias. """ def __init__( self, *args, min: FloatLike = 0.01, max: FloatLike = 1.0, **kwargs): super().__init__(*args, **kwargs, bias=False) self.min = min self.max = max # initialize weight to all positive values. with torch.no_grad(): self.weight[:] = 1.0 / self.weight[0][0].numel() def forward(self, input: Tensor) -> Tensor: """ Forward function. Input and returned tensor have shape: (N, C, H) i.e. (batch_size, num_channels, height) """ weight = limit_param_value(self.weight, min=float(self.min), max=float(self.max), training=self.training) # make absolutely sure there are no negative values. For parameter-averaging-related # reasons, we prefer to also use limit_param_value to make sure the weights stay # positive. weight = weight.abs() if self.padding_mode != 'zeros': return F.conv1d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode), weight, self.bias, self.stride, _single(0), self.dilation, self.groups) return F.conv1d(input, weight, self.bias, self.stride, self.padding, self.dilation, self.groups) class ConvNorm1d(torch.nn.Module): """ This is like BasicNorm except the denominator is summed over time using convolution with positive weights. Args: num_channels: the number of channels, e.g. 512. eps: the initial "epsilon" that we add as ballast in: scale = ((input_vec**2).mean() + epsilon)**-0.5 Note: our epsilon is actually large, but we keep the name to indicate the connection with conventional LayerNorm. learn_eps: if true, we learn epsilon; if false, we keep it at the initial value. eps_min: float eps_max: float """ def __init__( self, num_channels: int, eps: float = 0.25, learn_eps: bool = True, eps_min: float = -3.0, eps_max: float = 3.0, conv_min: float = 0.001, conv_max: float = 1.0, kernel_size: int = 15, ) -> None: super().__init__() self.num_channels = num_channels if learn_eps: self.eps = nn.Parameter(torch.tensor(eps).log().detach()) else: self.register_buffer("eps", torch.tensor(eps).log().detach()) self.eps_min = eps_min self.eps_max = eps_max pad = kernel_size // 2 # it has bias=False. self.conv = PositiveConv1d(1, 1, kernel_size=kernel_size, padding=pad, min=conv_min, max=conv_max) def forward(self, x: Tensor, src_key_padding_mask: Optional[Tensor] = None) -> Tensor: """ x shape: (N, C, T) src_key_padding_mask: the mask for the src keys per batch (optional): (N, T), contains True in masked positions. """ assert x.ndim == 3 and x.shape[1] == self.num_channels eps = self.eps if self.training and random.random() < 0.25: # with probability 0.25, in training mode, clamp eps between the min # and max; this will encourage it to learn parameters within the # allowed range by making parameters that are outside the allowed # range noisy. # gradients to allow the parameter to get back into the allowed # region if it happens to exit it. eps = torch.clamp(eps, min=self.eps_min, max=self.eps_max) # sqnorms: (N, 1, T) sqnorms = ( torch.mean(x ** 2, dim=1, keepdim=True) ) # 'counts' is a mechanism to correct for edge effects. counts = torch.ones_like(sqnorms) if src_key_padding_mask is not None: counts = counts.masked_fill_(src_key_padding_mask.unsqueeze(1), 0.0) sqnorms = sqnorms * counts sqnorms = self.conv(sqnorms) # the clamping is to avoid division by zero for padding frames. counts = torch.clamp(self.conv(counts), min=0.01) # scales: (N, 1, T) scales = (sqnorms / counts + eps.exp()) ** -0.5 # return x * scales class PositiveConv2d(nn.Conv2d): """ A modified form of nn.Conv2d where the weight parameters are constrained to be positive and there is no bias. """ def __init__( self, *args, min: FloatLike = 0.01, max: FloatLike = 1.0, **kwargs): super().__init__(*args, **kwargs, bias=False) self.min = min self.max = max # initialize weight to all positive values. with torch.no_grad(): self.weight[:] = 1.0 / self.weight[0][0].numel() def forward(self, input: Tensor) -> Tensor: """ Forward function. Input and returned tensor have shape: (N, C, H, W) i.e. (batch_size, num_channels, height, width) """ weight = limit_param_value(self.weight, min=float(self.min), max=float(self.max), training=self.training) # make absolutely sure there are no negative values. For parameter-averaging-related # reasons, we prefer to also use limit_param_value to make sure the weights stay # positive. weight = weight.abs() if self.padding_mode != 'zeros': return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode), weight, self.bias, self.stride, _pair(0), self.dilation, self.groups) return F.conv2d(input, weight, self.bias, self.stride, self.padding, self.dilation, self.groups) class ConvNorm2d(torch.nn.Module): """ This is like BasicNorm except the denominator is summed over time using convolution with positive weights. Args: num_channels: the number of channels, e.g. 512. eps: the initial "epsilon" that we add as ballast in: scale = ((input_vec**2).mean() + epsilon)**-0.5 Note: our epsilon is actually large, but we keep the name to indicate the connection with conventional LayerNorm. learn_eps: if true, we learn epsilon; if false, we keep it at the initial value. eps_min: float eps_max: float """ def __init__( self, num_channels: int, eps: float = 0.25, learn_eps: bool = True, eps_min: float = -3.0, eps_max: float = 3.0, conv_min: float = 0.001, conv_max: float = 1.0, kernel_size: Tuple[int, int] = (3, 3), ) -> None: super().__init__() self.num_channels = num_channels if learn_eps: self.eps = nn.Parameter(torch.tensor(eps).log().detach()) else: self.register_buffer("eps", torch.tensor(eps).log().detach()) self.eps_min = eps_min self.eps_max = eps_max pad = (kernel_size[0] // 2, kernel_size[1] // 2) # it has bias=False. self.conv = PositiveConv2d(1, 1, kernel_size=kernel_size, padding=pad, min=conv_min, max=conv_max) def forward(self, x: Tensor) -> Tensor: """ x shape: (N, C, H, W) """ assert x.ndim == 4 and x.shape[1] == self.num_channels eps = self.eps if self.training and random.random() < 0.25: # with probability 0.25, in training mode, clamp eps between the min # and max; this will encourage it to learn parameters within the # allowed range by making parameters that are outside the allowed # range noisy. # gradients to allow the parameter to get back into the allowed # region if it happens to exit it. eps = torch.clamp(eps, min=self.eps_min, max=self.eps_max) # sqnorms: (N, 1, H, W) sqnorms = ( torch.mean(x ** 2, dim=1, keepdim=True) ) # 'counts' is a mechanism to correct for edge effects. # TODO: key-padding mask counts = torch.ones_like(sqnorms) #if src_key_padding_mask is not None: # counts = counts.masked_fill_(src_key_padding_mask.unsqueeze(1), 0.0) #sqnorms = sqnorms * counts sqnorms = self.conv(sqnorms) # the clamping is to avoid division by zero for padding frames. counts = torch.clamp(self.conv(counts), min=0.01) # scales: (N, 1, H, W) scales = (sqnorms / counts + eps.exp()) ** -0.5 return x * scales def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs ) -> nn.Linear: """ Behaves like a constructor of a modified version of nn.Linear that gives an easy way to set the default initial parameter scale. Args: Accepts the standard args and kwargs that nn.Linear accepts e.g. in_features, out_features, bias=False. initial_scale: you can override this if you want to increase or decrease the initial magnitude of the module's output (affects the initialization of weight_scale and bias_scale). Another option, if you want to do something like this, is to re-initialize the parameters. """ ans = nn.Linear(*args, **kwargs) with torch.no_grad(): ans.weight[:] *= initial_scale if ans.bias is not None: torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale) return ans def ScaledConv1d(*args, initial_scale: float = 1.0, **kwargs ) -> nn.Conv1d: """ Behaves like a constructor of a modified version of nn.Conv1d that gives an easy way to set the default initial parameter scale. Args: Accepts the standard args and kwargs that nn.Linear accepts e.g. in_features, out_features, bias=False. initial_scale: you can override this if you want to increase or decrease the initial magnitude of the module's output (affects the initialization of weight_scale and bias_scale). Another option, if you want to do something like this, is to re-initialize the parameters. """ ans = nn.Conv1d(*args, **kwargs) with torch.no_grad(): ans.weight[:] *= initial_scale if ans.bias is not None: torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale) return ans def ScaledConv2d(*args, initial_scale: float = 1.0, **kwargs ) -> nn.Conv2d: """ Behaves like a constructor of a modified version of nn.Conv1d that gives an easy way to set the default initial parameter scale. Args: Accepts the standard args and kwargs that nn.Linear accepts e.g. in_features, out_features, bias=False. initial_scale: you can override this if you want to increase or decrease the initial magnitude of the module's output (affects the initialization of weight_scale and bias_scale). Another option, if you want to do something like this, is to re-initialize the parameters. """ ans = nn.Conv2d(*args, **kwargs) with torch.no_grad(): ans.weight[:] *= initial_scale if ans.bias is not None: torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale) return ans class ActivationBalancer(torch.nn.Module): """ Modifies the backpropped derivatives of a function to try to encourage, for each channel, that it is positive at least a proportion `threshold` of the time. It does this by multiplying negative derivative values by up to (1+max_factor), and positive derivative values by up to (1-max_factor), interpolated from 1 at the threshold to those extremal values when none of the inputs are positive. Args: num_channels: the number of channels channel_dim: the dimension/axis corresponding to the channel, e.g. -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative. min_positive: the minimum, per channel, of the proportion of the time that (x > 0), below which we start to modify the derivatives. max_positive: the maximum, per channel, of the proportion of the time that (x > 0), above which we start to modify the derivatives. max_factor: the maximum factor by which we modify the derivatives for either the sign constraint or the magnitude constraint; e.g. with max_factor=0.02, the the derivatives would be multiplied by values in the range [0.98..1.02]. sign_gain_factor: determines the 'gain' with which we increase the change in gradient once the constraints on min_positive and max_positive are violated. scale_gain_factor: determines the 'gain' with which we increase the change in gradient once the constraints on min_abs and max_abs are violated. min_abs: the minimum average-absolute-value difference from the mean value per channel, which we allow, before we start to modify the derivatives to prevent this. max_abs: the maximum average-absolute-value difference from the mean value per channel, which we allow, before we start to modify the derivatives to prevent this. prob: determines the minimum probability with which we modify the gradients for the {min,max}_positive and {min,max}_abs constraints, on each forward(). This is done randomly to prevent all layers from doing it at the same time. """ def __init__( self, num_channels: int, channel_dim: int, min_positive: FloatLike = 0.05, max_positive: FloatLike = 0.95, max_factor: FloatLike = 0.04, sign_gain_factor: FloatLike = 0.04, scale_gain_factor: FloatLike = 0.04, min_abs: FloatLike = 0.2, max_abs: FloatLike = 100.0, prob: Optional[FloatLike] = None, ): super(ActivationBalancer, self).__init__() if prob is None: prob = ScheduledFloat((0.0, 0.5), (8000.0, 0.125), default=0.4) self.prob = prob # 5% of the time we will return and do nothing because memory usage is # too high. self.mem_cutoff = CutoffEstimator(0.05) # actually self.num_channels is no longer needed except for an assertion. self.num_channels = num_channels self.channel_dim = channel_dim self.min_positive = min_positive self.max_positive = max_positive self.max_factor = max_factor self.min_abs = min_abs self.max_abs = max_abs self.sign_gain_factor = sign_gain_factor self.scale_gain_factor = scale_gain_factor def forward(self, x: Tensor) -> Tensor: if (torch.jit.is_scripting() or not x.requires_grad or (x.is_cuda and self.mem_cutoff(torch.cuda.memory_allocated()))): return _no_op(x) prob = float(self.prob) if random.random() < prob: assert x.shape[self.channel_dim] == self.num_channels sign_gain_factor = 0.5 if float(self.min_positive) != 0.0 or float(self.max_positive) != 1.0: sign_factor = _compute_sign_factor(x.detach(), self.channel_dim, float(self.min_positive), float(self.max_positive), gain_factor=float(self.sign_gain_factor) / prob, max_factor=float(self.max_factor)) else: sign_factor = None scale_factor, mean = _compute_scale_factor(x.detach(), self.channel_dim, min_abs=float(self.min_abs), max_abs=float(self.max_abs), gain_factor=float(self.scale_gain_factor) / prob, max_factor=float(self.max_factor)) return ActivationBalancerFunction.apply( x, scale_factor, mean, sign_factor, self.channel_dim, ) else: return _no_op(x) class BalancerFunction(torch.autograd.Function): @staticmethod def forward( ctx, x: Tensor, min_mean: float, max_mean: float, min_rms: float, max_rms: float, grad_scale: float, channel_dim: int, ) -> Tensor: if channel_dim < 0: channel_dim += x.ndim ctx.channel_dim = channel_dim ctx.save_for_backward(x) ctx.config = (min_mean, max_mean, min_rms, max_rms, grad_scale, channel_dim) return x @staticmethod def backward( ctx, x_grad: Tensor ) -> Tuple[Tensor, None, None, None, None, None]: x, = ctx.saved_tensors (min_mean, max_mean, min_rms, max_rms, grad_scale, channel_dim) = ctx.config with torch.enable_grad(): with torch.cuda.amp.autocast(enabled=False): x = x.to(torch.float32) x = x.detach() x.requires_grad = True mean_dims = [ i for i in range(x.ndim) if i != channel_dim ] uncentered_var = (x ** 2).mean(dim=mean_dims, keepdim=True) mean = x.mean(dim=mean_dims, keepdim=True) stddev = (uncentered_var - (mean * mean)).clamp(min=1.0e-20).sqrt() rms = uncentered_var.clamp(min=1.0e-20).sqrt() m = mean / stddev # part of loss that relates to mean / stddev m_loss = (m - m.clamp(min=min_mean, max=max_mean)).abs() # put a much larger scale on the RMS-max-limit loss, so that if both it and the # m_loss are violated we fix the RMS loss first. rms_clamped = rms.clamp(min=min_rms, max=max_rms) r_loss = (rms_clamped / rms).log().abs() loss = (m_loss + r_loss) loss.backward(gradient=torch.ones_like(loss)) loss_grad = x.grad loss_grad_rms = (loss_grad ** 2).mean(dim=mean_dims, keepdim=True).sqrt().clamp(min=1.0e-20) loss_grad = loss_grad * (grad_scale / loss_grad_rms) x_grad_float = x_grad.to(torch.float32) # scale each element of loss_grad by the absolute value of the corresponding # element of x_grad, which we view as a noisy estimate of its magnitude for that # (frame and dimension). later we can consider factored versions. x_grad_mod = x_grad_float + (x_grad_float.abs() * loss_grad) x_grad = x_grad_mod.to(x_grad.dtype) return x_grad, None, None, None, None, None, None class Balancer(torch.nn.Module): """ Modifies the backpropped derivatives of a function to try to encourage, for each channel, that it is positive at least a proportion `threshold` of the time. It does this by multiplying negative derivative values by up to (1+max_factor), and positive derivative values by up to (1-max_factor), interpolated from 1 at the threshold to those extremal values when none of the inputs are positive. Args: num_channels: the number of channels channel_dim: the dimension/axis corresponding to the channel, e.g. -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative. min_positive: the minimum, per channel, of the proportion of the time that (x > 0), below which we start to modify the derivatives. max_positive: the maximum, per channel, of the proportion of the time that (x > 0), above which we start to modify the derivatives. scale_gain_factor: determines the 'gain' with which we increase the change in gradient once the constraints on min_abs and max_abs are violated. min_abs: the minimum average-absolute-value difference from the mean value per channel, which we allow, before we start to modify the derivatives to prevent this. max_abs: the maximum average-absolute-value difference from the mean value per channel, which we allow, before we start to modify the derivatives to prevent this. prob: determines the minimum probability with which we modify the gradients for the {min,max}_positive and {min,max}_abs constraints, on each forward(). This is done randomly to prevent all layers from doing it at the same time. """ def __init__( self, num_channels: int, channel_dim: int, min_positive: FloatLike = 0.05, max_positive: FloatLike = 0.95, min_abs: FloatLike = 0.2, max_abs: FloatLike = 100.0, grad_scale: FloatLike = 0.04, prob: Optional[FloatLike] = None, ): super().__init__() if prob is None: prob = ScheduledFloat((0.0, 0.5), (8000.0, 0.125), default=0.4) self.prob = prob # 5% of the time we will return and do nothing because memory usage is # too high. self.mem_cutoff = CutoffEstimator(0.05) # actually self.num_channels is no longer needed except for an assertion. self.num_channels = num_channels self.channel_dim = channel_dim self.min_positive = min_positive self.max_positive = max_positive self.min_abs = min_abs self.max_abs = max_abs self.grad_scale = grad_scale def forward(self, x: Tensor) -> Tensor: if (torch.jit.is_scripting() or not x.requires_grad or (x.is_cuda and self.mem_cutoff(torch.cuda.memory_allocated()))): return _no_op(x) prob = float(self.prob) if random.random() < prob: # The following inner-functions convert from the way we historically specified # these limitations, as limits on the absolute value and the proportion of positive # values, to limits on the RMS value and the (mean / stddev). def _abs_to_rms(x): # for normally distributed data, if the expected absolute value is x, the # expected rms value will be sqrt(pi/2) * x. return 1.25331413732 * x def _proportion_positive_to_mean(x): def _atanh(x): eps = 1.0e-10 # eps is to prevent crashes if x is exactly 0 or 1. # we'll just end up returning a fairly large value. return (math.log (1+x+eps) - math.log (1-x+eps)) / 2. def _approx_inverse_erf(x): # 1 / (sqrt(pi) * ln(2)), # see https://math.stackexchange.com/questions/321569/approximating-the-error-function-erf-by-analytical-functions # this approximation is extremely crude and gets progressively worse for # x very close to -1 or +1, but we mostly care about the "middle" region # e.g. _approx_inverse_erf(0.05) = 0.0407316414078772, # and math.erf(0.0407316414078772) = 0.045935330944660666, # which is pretty close to 0.05. return 0.8139535143 * _atanh(x) # first convert x from the range 0..1 to the range -1..1 which the error # function returns x = -1 + (2 * x) return _approx_inverse_erf(x) min_mean = _proportion_positive_to_mean(float(self.min_positive)) max_mean = _proportion_positive_to_mean(float(self.max_positive)) min_rms = _abs_to_rms(float(self.min_abs)) max_rms = _abs_to_rms(float(self.max_abs)) grad_scale = float(self.grad_scale) assert x.shape[self.channel_dim] == self.num_channels return BalancerFunction.apply( x, min_mean, max_mean, min_rms, max_rms, grad_scale, self.channel_dim ) else: return _no_op(x) def penalize_abs_values_gt(x: Tensor, limit: float, penalty: float, name: str = None) -> Tensor: """ Returns x unmodified, but in backprop will put a penalty for the excess of the absolute values of elements of x over the limit "limit". E.g. if limit == 10.0, then if x has any values over 10 it will get a penalty. Caution: the value of this penalty will be affected by grad scaling used in automatic mixed precision training. For this reasons we use this, it shouldn't really matter, or may even be helpful; we just use this to disallow really implausible values of scores to be given to softmax. The name is for randomly printed debug info. """ x_sign = x.sign() over_limit = (x.abs() - limit) > 0 # The following is a memory efficient way to penalize the absolute values of # x that's over the limit. (The memory efficiency comes when you think # about which items torch needs to cache for the autograd, and which ones it # can throw away). The numerical value of aux_loss as computed here will # actually be larger than it should be, by limit * over_limit.sum(), but it # has the same derivative as the real aux_loss which is penalty * (x.abs() - # limit).relu(). aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x) # note: we don't do sum() here on aux)_loss, but it's as if we had done # sum() due to how with_loss() works. x = with_loss(x, aux_loss, name) # you must use x for something, or this will be ineffective. return x def _diag(x: Tensor): # like .diag(), but works for tensors with 3 dims. if x.ndim == 2: return x.diag() else: (batch, dim, dim) = x.shape x = x.reshape(batch, dim * dim) x = x[:, ::dim+1] assert x.shape == (batch, dim) return x def _whitening_metric(x: Tensor, num_groups: int): """ Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of of the centered feature covariance are the same within each group's covariance matrix and also between groups. Args: x: a Tensor of shape (*, num_channels) num_groups: the number of groups of channels, a number >=1 that divides num_channels Returns: Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and greater than 1.0 otherwise. """ assert x.dtype != torch.float16 x = x.reshape(-1, x.shape[-1]) (num_frames, num_channels) = x.shape assert num_channels % num_groups == 0 channels_per_group = num_channels // num_groups x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1) # x now has shape (num_groups, num_frames, channels_per_group) # subtract the mean so we use the centered, not uncentered, covariance. # My experience has been that when we "mess with the gradients" like this, # it's better not do anything that tries to move the mean around, because # that can easily cause instability. x = x - x.mean(dim=1, keepdim=True) # x_covar: (num_groups, channels_per_group, channels_per_group) x_covar = torch.matmul(x.transpose(1, 2), x) x_covar_mean_diag = _diag(x_covar).mean() # the following expression is what we'd get if we took the matrix product # of each covariance and measured the mean of its trace, i.e. # the same as _diag(torch.matmul(x_covar, x_covar)).mean(). x_covarsq_mean_diag = (x_covar ** 2).sum() / (num_groups * channels_per_group) # this metric will be >= 1.0; the larger it is, the less 'white' the data was. metric = x_covarsq_mean_diag / (x_covar_mean_diag ** 2 + 1.0e-20) return metric class WhiteningPenaltyFunction(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, module: nn.Module) -> Tensor: ctx.save_for_backward(x) ctx.module = module return x @staticmethod def backward(ctx, x_grad: Tensor): x_orig, = ctx.saved_tensors w = ctx.module with torch.enable_grad(): with torch.cuda.amp.autocast(enabled=False): x_detached = x_orig.to(torch.float32).detach() x_detached.requires_grad = True metric = _whitening_metric(x_detached, w.num_groups) if random.random() < 0.005 or __name__ == "__main__": logging.info(f"Whitening: name={w.name}, num_groups={w.num_groups}, num_channels={x_orig.shape[-1]}, " f"metric={metric.item():.2f} vs. limit={float(w.whitening_limit)}") if metric < float(w.whitening_limit): w.prob = w.min_prob return x_grad, None else: w.prob = w.max_prob metric.backward() penalty_grad = x_detached.grad scale = w.grad_scale * (x_grad.to(torch.float32).norm() / (penalty_grad.norm() + 1.0e-20)) penalty_grad = penalty_grad * scale return x_grad + penalty_grad.to(x_grad.dtype), None class Whiten(nn.Module): def __init__( self, num_groups: int, whitening_limit: FloatLike, prob: Union[float, Tuple[float,float]], grad_scale: FloatLike): """ Args: num_groups: the number of groups to divide the channel dim into before whitening. We will attempt to make the feature covariance within each group, after mean subtraction, as "white" as possible, while having the same trace across all groups. whitening_limit: a value greater than 1.0, that dictates how much freedom we have to violate the constraints. 1.0 would mean perfectly white, with exactly the same trace across groups; larger values give more freedom. E.g. 2.0. prob: the probability with which we apply the gradient modification (also affects the grad scale). May be supplied as a float, or as a pair (min_prob, max_prob) grad_scale: determines the scale on the gradient term from this object, relative to the rest of the gradient on the attention weights. E.g. 0.02 (you may want to use smaller values than this if prob is large) """ super(Whiten, self).__init__() assert num_groups >= 1 assert float(whitening_limit) >= 1 assert grad_scale >= 0 self.num_groups = num_groups self.whitening_limit = whitening_limit self.grad_scale = grad_scale if isinstance(prob, float): prob = (prob, prob) (self.min_prob, self.max_prob) = prob assert 0 < self.min_prob <= self.max_prob <= 1 self.prob = self.max_prob self.name = None # will be set in training loop def forward(self, x: Tensor) -> Tensor: """ In the forward pass, this function just returns the input unmodified. In the backward pass, it will modify the gradients to ensure that the distribution in each group has close to (lambda times I) as the covariance after mean subtraction, with the same lambda across groups. For whitening_limit > 1, there will be more freedom to violate this constraint. Args: x: the input of shape (*, num_channels) Returns: x, unmodified. You should make sure you use the returned value, or the graph will be freed and nothing will happen in backprop. """ grad_scale = float(self.grad_scale) if not x.requires_grad or random.random() > self.prob or grad_scale == 0: return _no_op(x) else: return WhiteningPenaltyFunction.apply(x, self) class WithLoss(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, y: Tensor, name: str): ctx.y_shape = y.shape if random.random() < 0.002 and name is not None: loss_sum = y.sum().item() logging.info(f"WithLoss: name={name}, loss-sum={loss_sum:.3e}") return x @staticmethod def backward(ctx, ans_grad: Tensor): return ans_grad, torch.ones(ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device), None def with_loss(x, y, name): # returns x but adds y.sum() to the loss function. return WithLoss.apply(x, y, name) class ScaleGradFunction(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, alpha: float) -> Tensor: ctx.alpha = alpha return x @staticmethod def backward(ctx, grad: Tensor): return grad * ctx.alpha, None def scale_grad(x: Tensor, alpha: float): return ScaleGradFunction.apply(x, alpha) class ScaleGrad(nn.Module): def __init__(self, alpha: float): super().__init__() self.alpha = alpha def forward(self, x: Tensor) -> Tensor: return scale_grad(x, self.alpha) class LimitParamValue(torch.autograd.Function): @staticmethod def forward(ctx, x: Tensor, min: float, max: float): ctx.save_for_backward(x) assert max >= min ctx.min = min ctx.max = max return x @staticmethod def backward(ctx, x_grad: Tensor): x, = ctx.saved_tensors # where x < ctx.min, ensure all grads are negative (this will tend to make # x more positive). x_grad = x_grad * torch.where(torch.logical_and(x_grad > 0, x < ctx.min), -1.0, 1.0) # where x > ctx.max, ensure all grads are positive (this will tend to make # x more negative). x_grad *= torch.where(torch.logical_and(x_grad < 0, x > ctx.max), -1.0, 1.0) return x_grad, None, None def limit_param_value(x: Tensor, min: float, max: float, prob: float = 0.6, training: bool = True): # You apply this to (typically) an nn.Parameter during training to ensure that its # (elements mostly) stays within a supplied range. This is done by modifying the # gradients in backprop. # It's not necessary to do this on every batch: do it only some of the time, # to save a little time. if training and random.random() < prob: return LimitParamValue.apply(x, min, max) else: return x def _no_op(x: Tensor) -> Tensor: if (torch.jit.is_scripting()): return x else: # a no-op function that will have a node in the autograd graph, # to avoid certain bugs relating to backward hooks return x.chunk(1, dim=-1)[0] class Identity(torch.nn.Module): def __init__(self): super(Identity, self).__init__() def forward(self, x): return _no_op(x) class MaxEig(torch.nn.Module): """ Modifies the backpropped derivatives of a function to try to discourage that any given direction in activation space accounts for more than a specified proportion of the covariance (e.g. 0.2). Args: num_channels: the number of channels channel_dim: the dimension/axis corresponding to the channel, e.g. -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative. max_var_per_eig: the maximum proportion of the variance of the features/channels, after mean subtraction, that can come from any given eigenvalue. min_prob: the minimum probability with which we apply this during any invocation of forward(), assuming last time we applied the constraint it was not active; supplied for speed. scale: determines the scale with which we modify the gradients, relative to the existing / unmodified gradients """ def __init__( self, num_channels: int, channel_dim: int, max_var_per_eig: float = 0.2, min_prob: float = 0.01, scale: float = 0.01, ): super(MaxEig, self).__init__() self.num_channels = num_channels self.channel_dim = channel_dim self.scale = scale assert max_var_per_eig == 0.0 or max_var_per_eig > 1.0 / num_channels self.max_var_per_eig = max_var_per_eig # we figure out the dominant direction using the power method: starting with # a random vector, keep multiplying by the covariance and renormalizing. with torch.no_grad(): # arbitrary.. would use randn() but want to leave the rest of the model's # random parameters unchanged for comparison direction = torch.arange(num_channels).to(torch.float) direction = direction / direction.norm() self.register_buffer('max_eig_direction', direction) self.min_prob = min_prob # cur_prob is the current probability we'll use to apply the ActivationBalancer. # We'll regress this towards prob, each time we try to apply it and it is not # active. self.cur_prob = 1.0 def forward(self, x: Tensor) -> Tensor: if (torch.jit.is_scripting() or self.max_var_per_eig <= 0 or random.random() > self.cur_prob): return _no_op(x) with torch.cuda.amp.autocast(enabled=False): eps = 1.0e-20 orig_x = x x = x.to(torch.float32) with torch.no_grad(): x = x.transpose(self.channel_dim, -1).reshape(-1, self.num_channels) x = x - x.mean(dim=0) new_direction, coeffs = self._find_direction_coeffs(x, self.max_eig_direction) x_var = (x**2).mean() x_residual = x - coeffs * new_direction x_residual_var = (x_residual**2).mean() # `variance_proportion` is the proportion of the variance accounted for # by the top eigen-direction. variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20) # ensure new direction is nonzero even if x == 0, by including `direction`. self._set_direction(0.1 * self.max_eig_direction + new_direction) if random.random() < 0.01 or __name__ == "__main__": logging.info(f"variance_proportion = {variance_proportion.item()}, shape={tuple(orig_x.shape)}, cur_prob={self.cur_prob}") if variance_proportion >= self.max_var_per_eig: # The constraint is active. Note, we should quite rarely # reach here, only near the beginning of training if we are # starting to diverge, should this constraint be active. cur_prob = self.cur_prob self.cur_prob = 1.0 # next time, do the update with probability 1.0. return MaxEigLimiterFunction.apply(orig_x, coeffs, new_direction, self.channel_dim, self.scale) else: # let self.cur_prob exponentially approach self.min_prob, as # long as the constraint is inactive. self.cur_prob = 0.75 * self.cur_prob + 0.25 * self.min_prob return orig_x def _set_direction(self, direction: Tensor): """ Sets self.max_eig_direction to a normalized version of `direction` """ direction = direction.detach() direction = direction / direction.norm() direction_sum = direction.sum().item() if direction_sum - direction_sum == 0: # no inf/nan self.max_eig_direction[:] = direction else: logging.info(f"Warning: sum of direction in MaxEig is {direction_sum}, " "num_channels={self.num_channels}, channel_dim={self.channel_dim}") def _find_direction_coeffs(self, x: Tensor, prev_direction: Tensor) -> Tuple[Tensor, Tensor]: """ Figure out (an approximation to) the proportion of the variance of a set of feature vectors that can be attributed to the top eigen-direction. Args: x: a Tensor of shape (num_frames, num_channels), with num_frames > 1. prev_direction: a Tensor of shape (num_channels,), that is our previous estimate of the top eigen-direction, or a random direction if this is the first iteration. Does not have to be normalized, but should be nonzero. Returns: (cur_direction, coeffs), where: cur_direction: a Tensor of shape (num_channels,) that is the current estimate of the top eigen-direction. coeffs: a Tensor of shape (num_frames, 1) that minimizes, or approximately minimizes, (x - coeffs * cur_direction).norm() """ (num_frames, num_channels) = x.shape assert num_channels > 1 and num_frames > 1 assert prev_direction.shape == (num_channels,) # `coeffs` are the coefficients of `prev_direction` in x. # actually represent the coeffs up to a constant positive factor. coeffs = (x * prev_direction).sum(dim=1, keepdim=True) + 1.0e-10 cur_direction = (x * coeffs).sum(dim=0) / ((coeffs ** 2).sum() + 1.0e-20) return cur_direction, coeffs class DoubleSwishFunction(torch.autograd.Function): """ double_swish(x) = x * torch.sigmoid(x-1) This is a definition, originally motivated by its close numerical similarity to swish(swish(x)), where swish(x) = x * sigmoid(x). Memory-efficient derivative computation: double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1) double_swish'(x) = d/dx double_swish(x) = x * s'(x) + x' * s(x) = x * s'(x) + s(x). Now, s'(x) = s(x) * (1-s(x)). double_swish'(x) = x * s'(x) + s(x). = x * s(x) * (1-s(x)) + s(x). = double_swish(x) * (1-s(x)) + s(x) ... so we just need to remember s(x) but not x itself. """ @staticmethod def forward(ctx, x: Tensor) -> Tensor: requires_grad = x.requires_grad x_dtype = x.dtype if x.dtype == torch.float16: x = x.to(torch.float32) s = torch.sigmoid(x - 1.0) y = x * s if requires_grad: deriv = (y * (1 - s) + s) # notes on derivative of x * sigmoid(x - 1): # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29 # min \simeq -0.043638. Take floor as -0.044 so it's a lower bund # max \simeq 1.1990. Take ceil to be 1.2 so it's an upper bound. # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which # floors), should be expectation-preserving. floor = -0.044 ceil = 1.2 d_scaled = ((deriv - floor) * (255.0 / (ceil - floor)) + torch.rand_like(deriv)) if __name__ == "__main__": # for self-testing only. assert d_scaled.min() >= 0.0 assert d_scaled.max() < 256.0 d_int = d_scaled.to(torch.uint8) ctx.save_for_backward(d_int) if x.dtype == torch.float16 or torch.is_autocast_enabled(): y = y.to(torch.float16) return y @staticmethod def backward(ctx, y_grad: Tensor) -> Tensor: d, = ctx.saved_tensors # the same constants as used in forward pass. floor = -0.043637 ceil = 1.2 d = (d * ((ceil - floor) / 255.0) + floor) return y_grad * d class DoubleSwish(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x: Tensor) -> Tensor: """Return double-swish activation function which is an approximation to Swish(Swish(x)), that we approximate closely with x * sigmoid(x-1). """ if torch.jit.is_scripting(): return x * torch.sigmoid(x - 1.0) return DoubleSwishFunction.apply(x) class TanSwishFunction(torch.autograd.Function): """ double_swish(x) = tan(x) * torch.sigmoid(x-1) entering: d/dx(tanh(x) * sigmoid(x-1)) into wolfram alpha, I see that the range of this function is -0.0498087 <= y <= 0.417894 let's make it (as we don't know how this was rounded): -0.0498088 <= y <= 0.417895 """ @staticmethod def forward(ctx, x: Tensor) -> Tensor: requires_grad = x.requires_grad if not requires_grad: return torch.tanh(x) * torch.sigmoid(x - 1.0) x_dtype = x.dtype if x.dtype == torch.float16: x = x.to(torch.float32) with torch.cuda.amp.autocast(enabled=False): with torch.enable_grad(): x = x.detach() x.requires_grad = True y = torch.tanh(x) * torch.sigmoid(x - 1.0) y.backward(gradient=torch.ones_like(y)) grad = x.grad floor = -0.0498088 ceil = 0.417895 d_scaled = ((grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(grad)) if __name__ == "__main__": # for self-testing only. assert d_scaled.min() >= 0.0 assert d_scaled.max() < 256.0 d_int = d_scaled.to(torch.uint8) ctx.save_for_backward(d_int) if x.dtype == torch.float16 or torch.is_autocast_enabled(): y = y.to(torch.float16) return y @staticmethod def backward(ctx, y_grad: Tensor) -> Tensor: d, = ctx.saved_tensors # the same constants as used in forward pass. floor = -0.0498088 ceil = 0.417895 d = (d * ((ceil - floor) / 255.0) + floor) return (y_grad * d) class TanSwish(torch.nn.Module): def forward(self, x: Tensor) -> Tensor: """Return tan-swish activation function which is tanh(x) sigmoid(x-1)n """ if torch.jit.is_scripting(): return x.tanh() * torch.sigmoid(x - 1.0) return TanSwishFunction.apply(x) # Dropout2 is just like normal dropout, except it supports schedules on the dropout rates. class Dropout2(nn.Module): def __init__(self, p: FloatLike): super().__init__() self.p = p def forward(self, x: Tensor) -> Tensor: return torch.nn.functional.dropout(x, p=float(self.p), training=self.training) class SwooshLFunction(torch.autograd.Function): """ swoosh(x) = log(1 + exp(x-4)) - 0.08*x - 0.035 """ @staticmethod def forward(ctx, x: Tensor) -> Tensor: requires_grad = x.requires_grad x_dtype = x.dtype if x.dtype == torch.float16: x = x.to(torch.float32) zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) coeff = -0.08 with torch.cuda.amp.autocast(enabled=False): with torch.enable_grad(): x = x.detach() x.requires_grad = True y = torch.logaddexp(zero, x - 4.0) + coeff * x - 0.035 if not requires_grad: return y y.backward(gradient = torch.ones_like(y)) grad = x.grad floor = coeff ceil = 1.0 + coeff + 0.005 d_scaled = ((grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(grad)) if __name__ == "__main__": # for self-testing only. assert d_scaled.min() >= 0.0 assert d_scaled.max() < 256.0 d_int = d_scaled.to(torch.uint8) ctx.save_for_backward(d_int) if x.dtype == torch.float16 or torch.is_autocast_enabled(): y = y.to(torch.float16) return y @staticmethod def backward(ctx, y_grad: Tensor) -> Tensor: d, = ctx.saved_tensors # the same constants as used in forward pass. coeff = -0.08 floor = coeff ceil = 1.0 + coeff + 0.005 d = (d * ((ceil - floor) / 255.0) + floor) return (y_grad * d) class SwooshL(torch.nn.Module): def forward(self, x: Tensor) -> Tensor: """Return Swoosh-L activation. """ if torch.jit.is_scripting(): zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) return torch.logaddexp(zero, x - 4.0) - 0.08 * x - 0.035 return SwooshLFunction.apply(x) class SwooshRFunction(torch.autograd.Function): """ swoosh(x) = log(1 + exp(x-1)) - 0.08*x - 0.313261687 derivatives are between -0.08 and 0.92. """ @staticmethod def forward(ctx, x: Tensor) -> Tensor: requires_grad = x.requires_grad x_dtype = x.dtype if x.dtype == torch.float16: x = x.to(torch.float32) zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) with torch.cuda.amp.autocast(enabled=False): with torch.enable_grad(): x = x.detach() x.requires_grad = True y = torch.logaddexp(zero, x - 1.) - 0.08 * x - 0.313261687 if not requires_grad: return y y.backward(gradient = torch.ones_like(y)) grad = x.grad floor = -0.08 ceil = 0.925 d_scaled = ((grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(grad)) if __name__ == "__main__": # for self-testing only. assert d_scaled.min() >= 0.0 assert d_scaled.max() < 256.0 d_int = d_scaled.to(torch.uint8) ctx.save_for_backward(d_int) if x.dtype == torch.float16 or torch.is_autocast_enabled(): y = y.to(torch.float16) return y @staticmethod def backward(ctx, y_grad: Tensor) -> Tensor: d, = ctx.saved_tensors # the same constants as used in forward pass. floor = -0.08 ceil = 0.925 d = (d * ((ceil - floor) / 255.0) + floor) return (y_grad * d) class SwooshR(torch.nn.Module): def forward(self, x: Tensor) -> Tensor: """Return Swoosh-L activation. """ if torch.jit.is_scripting(): zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) return torch.logaddexp(zero, x - 1.) - 0.08 * x - 0.313261687 return SwooshRFunction.apply(x) def _test_max_eig(): for proportion in [0.1, 0.5, 10.0]: logging.info(f"proportion = {proportion}") x = torch.randn(100, 128) direction = torch.randn(128) coeffs = torch.randn(100, 1) x += proportion * direction * coeffs x.requires_grad = True num_channels = 128 m = MaxEig(num_channels, 1, # channel_dim 0.5, # max_var_per_eig scale=0.1) # grad_scale for _ in range(4): y = m(x) y_grad = torch.randn_like(x) y.backward(gradient=y_grad) if proportion < 0.2: assert torch.allclose(x.grad, y_grad, atol=1.0e-02) elif proportion > 1.0: assert not torch.allclose(x.grad, y_grad) def _test_whiten(): for proportion in [0.1, 0.5, 10.0]: logging.info(f"_test_whiten(): proportion = {proportion}") x = torch.randn(100, 128) direction = torch.randn(128) coeffs = torch.randn(100, 1) x += proportion * direction * coeffs x.requires_grad = True num_channels = 128 m = Whiten(1, # num_groups 5.0, # whitening_limit, prob=1.0, grad_scale=0.1) # grad_scale for _ in range(4): y = m(x) y_grad = torch.randn_like(x) y.backward(gradient=y_grad) if proportion < 0.2: assert torch.allclose(x.grad, y_grad) elif proportion > 1.0: assert not torch.allclose(x.grad, y_grad) def _test_activation_balancer_sign(): probs = torch.arange(0, 1, 0.01) N = 1000 x = 1.0 * ((2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0) x = x.detach() x.requires_grad = True m = ActivationBalancer( probs.numel(), channel_dim=0, min_positive=0.05, max_positive=0.95, max_factor=0.2, min_abs=0.0, prob=1.0, ) y_grad = torch.sign(torch.randn(probs.numel(), N)) y = m(x) y.backward(gradient=y_grad) print("_test_activation_balancer_sign: x = ", x) print("_test_activation_balancer_sign: y grad = ", y_grad) print("_test_activation_balancer_sign: x grad = ", x.grad) def _test_balancer_sign(): probs = torch.arange(0, 1, 0.01) N = 1000 x = 1.0 * ((2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0) x = x.detach() x.requires_grad = True m = Balancer( probs.numel(), channel_dim=0, min_positive=0.05, max_positive=0.95, min_abs=0.0, prob=1.0, ) y_grad = torch.sign(torch.randn(probs.numel(), N)) y = m(x) y.backward(gradient=y_grad) print("_test_balancer_sign: x = ", x) print("_test_balancer_sign: y grad = ", y_grad) print("_test_balancer_sign: x grad = ", x.grad) def _test_activation_balancer_magnitude(): magnitudes = torch.arange(0, 1, 0.01) N = 1000 x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze( -1 ) x = x.detach() x.requires_grad = True m = ActivationBalancer( magnitudes.numel(), channel_dim=0, min_positive=0.0, max_positive=1.0, max_factor=0.2, min_abs=0.2, max_abs=0.8, prob=1.0, ) y_grad = torch.sign(torch.randn(magnitudes.numel(), N)) y = m(x) y.backward(gradient=y_grad) print("_test_activation_balancer_magnitude: x = ", x) print("_test_activation_balancer_magnitude: y grad = ", y_grad) print("_test_activation_balancer_magnitude: x grad = ", x.grad) def _test_balancer_magnitude(): magnitudes = torch.arange(0, 1, 0.01) N = 1000 x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze( -1 ) x = x.detach() x.requires_grad = True m = Balancer( magnitudes.numel(), channel_dim=0, min_positive=0.0, max_positive=1.0, min_abs=0.2, max_abs=0.7, prob=1.0, ) y_grad = torch.sign(torch.randn(magnitudes.numel(), N)) y = m(x) y.backward(gradient=y_grad) print("_test_balancer_magnitude: x = ", x) print("_test_balancer_magnitude: y grad = ", y_grad) print("_test_balancer_magnitude: x grad = ", x.grad) def _test_basic_norm(): num_channels = 128 m = BasicNorm(num_channels=num_channels, channel_dim=1) x = torch.randn(500, num_channels) y = m(x) assert y.shape == x.shape x_rms = (x ** 2).mean().sqrt() y_rms = (y ** 2).mean().sqrt() print("x rms = ", x_rms) print("y rms = ", y_rms) def _test_double_swish_deriv(): x = torch.randn(10, 12, dtype=torch.double) * 3.0 x.requires_grad = True m = DoubleSwish() tol = ((1.2-(-0.043637))/255.0) torch.autograd.gradcheck(m, x, atol=tol) # for self-test. x = torch.randn(1000, 1000, dtype=torch.double) * 3.0 x.requires_grad = True y = m(x) def _test_tan_swish_deriv(): x = torch.randn(10, 12, dtype=torch.double) * 3.0 x.requires_grad = True m = TanSwish() tol = ((1.2-(-0.043637))/255.0) torch.autograd.gradcheck(m, x, atol=tol) # for self-test. x = torch.randn(1000, 1000, dtype=torch.double) * 3.0 x.requires_grad = True y = m(x) def _test_swooshl_deriv(): x = torch.randn(10, 12, dtype=torch.double) * 3.0 x.requires_grad = True m = SwooshL() tol = (1.0 / 255.0) torch.autograd.gradcheck(m, x, atol=tol) # for self-test. x = torch.randn(1000, 1000, dtype=torch.double) * 3.0 x.requires_grad = True y = m(x) def _test_swooshr_deriv(): x = torch.randn(10, 12, dtype=torch.double) * 3.0 x.requires_grad = True m = SwooshR() tol = (1.0 / 255.0) torch.autograd.gradcheck(m, x, atol=tol) # for self-test. x = torch.randn(1000, 1000, dtype=torch.double) * 3.0 x.requires_grad = True y = m(x) def _test_softmax(): a = torch.randn(2, 10, dtype=torch.float64) b = a.clone() a.requires_grad = True b.requires_grad = True a.softmax(dim=1)[:,0].sum().backward() print("a grad = ", a.grad) softmax(b, dim=1)[:,0].sum().backward() print("b grad = ", b.grad) assert torch.allclose(a.grad, b.grad) if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) torch.set_num_threads(1) torch.set_num_interop_threads(1) _test_softmax() _test_whiten() _test_max_eig() _test_activation_balancer_sign() _test_balancer_sign() _test_activation_balancer_magnitude() _test_balancer_magnitude() _test_basic_norm() _test_double_swish_deriv() _test_tan_swish_deriv() _test_swooshr_deriv() _test_swooshl_deriv()