1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124 | import torch
class BatchRenorm(torch.nn.Module):
"""BatchRenorm Module (https://arxiv.org/abs/1702.03275).
Adapted to Pytorch from
https://github.com/araffin/sbx/blob/master/sbx/common/jax_layers.py.
BatchRenorm is an improved version of vanilla BatchNorm. Contrary to BatchNorm,
BatchRenorm uses the running statistics for normalizing the batches after a warmup phase.
This makes it less prone to suffer from "outlier" batches that can happen
during very long training runs and, therefore, is more robust during long training runs.
During the warmup phase, it behaves exactly like a BatchNorm layer. After the warmup phase,
the running statistics are used for normalization. The running statistics are updated during
training mode. During evaluation mode, the running statistics are used for normalization but
not updated.
Args:
num_features: Number of features in the input tensor.
eps: A value added to the variance for numerical stability.
momentum: The value used for the ra_mean and ra_var (running average) computation.
It controls the rate of convergence for the batch renormalization statistics.
affine: A boolean value that when set to True, this module has learnable
affine parameters. Default: True
warmup_steps: Number of warum steps that are performed before the running statistics
are used for normalization. During the warump phase, the batch statistics are used.
"""
def __init__(
self,
num_features: int,
eps: float = 0.001,
momentum: float = 0.01,
affine: bool = True,
warmup_steps: int = 100_000,
) -> None:
super().__init__()
# Running average mean and variance
self.register_buffer("ra_mean", torch.zeros(num_features, dtype=torch.float))
self.register_buffer("ra_var", torch.ones(num_features, dtype=torch.float))
self.register_buffer("steps", torch.tensor(0, dtype=torch.long))
self.scale = torch.nn.Parameter(torch.ones(num_features, dtype=torch.float))
self.bias = torch.nn.Parameter(torch.zeros(num_features, dtype=torch.float))
self.affine = affine
self.eps = eps
self.momentum = momentum
self.num_features = num_features
# Clip scale and bias of the affine transform
self.rmax = 3.0
self.dmax = 5.0
self.warmup_steps = warmup_steps
def _check_input_dim(self, x: torch.Tensor) -> None:
raise NotImplementedError()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Normalize the input tensor.
Args:
x: Input tensor
Returns:
Normalized tensor.
"""
if self.training:
# Compute batch statistics
batch_mean = x.mean(dim=0)
batch_var = x.var(dim=0)
batch_std = (batch_var + self.eps).sqrt()
# Use batch statistics during initial warm up phase.
# Note: in the original paper, after some warmup phase (batch norm phase of 5k steps)
# the constraints are linearly relaxed to r_max/d_max over 40k steps
# Here we only have a warmup phase
if self.steps > self.warmup_steps:
running_std = (self.ra_var + self.eps).sqrt()
# scale
r = (batch_std / running_std).detach()
r = r.clamp(1 / self.rmax, self.rmax)
# bias
d = ((batch_mean - self.ra_mean) / running_std).detach()
d = d.clamp(-self.dmax, self.dmax)
# BatchNorm normalization, using minibatch stats and running average stats
custom_mean = batch_mean - d * batch_var.sqrt() / r
custom_var = batch_var / (r**2)
else:
custom_mean, custom_var = batch_mean, batch_var
# Update Running Statistics
self.ra_mean += self.momentum * (batch_mean.detach() - self.ra_mean)
self.ra_var += self.momentum * (batch_var.detach() - self.ra_var)
self.steps += 1
else:
# Use running statistics during evaluation mode
custom_mean, custom_var = self.ra_mean, self.ra_var
# Normalize
x = (x - custom_mean[None]) / (custom_var[None] + self.eps).sqrt()
if self.affine:
x = self.scale * x + self.bias
return x
def extra_repr(self) -> str:
return (
f"num_features={self.num_features}, momentum={self.momentum}, "
f"warmup_steps={self.warmup_steps}, affine={self.affine}"
)
class BatchRenorm1d(BatchRenorm):
def _check_input_dim(self, x: torch.Tensor) -> None:
if x.dim() == 1:
raise ValueError(f"Expected 2D or 3D input (got {x.dim()}D input)")
__all__ = ["BatchRenorm", "BatchRenorm1d"]
|