Skip to content

Improve numerical stability of batch normalization on CPUs #3113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
davisking merged 4 commits into from
Sep 29, 2025
+48 −46
Merged

Improve numerical stability of batch normalization on CPUs #3113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3dd187c
Make test cases `test_batch_normalize` and `test_batch_normalize_conv…
reunanen Sep 1, 2025
25bcd14
Improve numerical stability of `cpu::batch_normalize` and `cpu::batch…
reunanen Sep 1, 2025
2f7dc21
Replace FindCUDA with FindCUDAToolkit (fix #2833) (#3090)
smu-sc-gj Sep 6, 2025
f564cc1
Merge branch 'refs/heads/master' into reunanen-improve-numerical-stab…
davisking Sep 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 46 additions & 44 deletions dlib/cuda/cpu_dlib.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -816,37 +816,38 @@ namespace dlib
dest.copy_size(src);
means.set_size(1, src.k(), src.nr(), src.nc());
invstds.set_size(1, src.k(), src.nr(), src.nc());
running_means.set_size(1, src.k(), src.nr(), src.nc());
running_variances.set_size(1, src.k(), src.nr(), src.nc());

// first compute means and invstds
means = 0;
invstds = 0;
const auto p_invstds = invstds.host();
const auto p_means = means.host();
auto p_src = src.host();
const auto rvar = running_variances.host();
const long num = src.k()*src.nr()*src.nc();
// compute means, and sum of squares

// This scale makes the running variances unbiased.
const double scale = (src.num_samples())/(src.num_samples()-1.0);

// Apply Welford's algorithm to improve numerical stability
for (long i = 0; i < num; ++i)
{
double mean = 0.0;
double M2 = 0.0;

for (long n = 0; n < src.num_samples(); ++n)
{
float val = p_src[n*num+i];
p_means[i] += val;
p_invstds[i] += val*val;
const double delta1 = val - mean;
mean += delta1 / (n + 1);
const double delta2 = val - mean;
M2 += delta1 * delta2;
}
}
means /= src.num_samples();
invstds /= src.num_samples();
// copy data back to host
invstds.host(); means.host();

// compute variances
running_variances.copy_size(invstds);
auto rvar = running_variances.host();
// This scale makes the running variances unbiased.
const double scale = (src.num_samples())/(src.num_samples()-1.0);
for (long i = 0; i < num; ++i)
{
auto actual_var = p_invstds[i] - p_means[i]*p_means[i];
p_means[i] = mean;

const auto actual_var = (src.num_samples() > 1) ? (M2 / src.num_samples()) : 0.0;

if (averaging_factor == 1)
rvar[i] = scale*actual_var;
else
Expand All @@ -855,7 +856,6 @@ namespace dlib
p_invstds[i] = 1.0f/std::sqrt(actual_var + eps);
}

p_src = src.host();
auto p_dest = dest.host();
const auto p_gamma = gamma.host();
const auto p_beta = beta.host();
Expand All @@ -871,7 +871,6 @@ namespace dlib
}

// now keep track of the running means
running_means.copy_size(means);
if (averaging_factor != 1)
running_means = (1-averaging_factor)*mat(running_means) + averaging_factor*mat(means);
else
Expand Down Expand Up @@ -1083,52 +1082,56 @@ namespace dlib
dest.copy_size(src);
means.set_size(1, src.k());
invstds.set_size(1, src.k());
running_means.set_size(1, src.k());
running_variances.set_size(1, src.k());

// first compute means and invstds
means = 0;
invstds = 0;
const auto p_invstds = invstds.host();
const auto p_means = means.host();
const auto p_gamma = gamma.host();
const auto p_beta = beta.host();
auto p_src = src.host();
auto rvar = running_variances.host();
const long num = src.nr()*src.nc();
// compute means, and sum of squares
for (long n = 0; n < src.num_samples(); ++n)

// This scale makes the running variances unbiased.
const double scale = (src.num_samples()*num)/(src.num_samples()*num-1.0);

// Apply Welford's algorithm to improve numerical stability
for (long k = 0; k < src.k(); ++k)
{
for (long k = 0; k < src.k(); ++k)
double mean = 0.0;
double M2 = 0.0;
long count = 0;

for (long n = 0; n < src.num_samples(); ++n)
{
long start_index = tensor_index(src, n, k, 0, 0);
auto p = p_src + start_index;

for (long i = 0; i < num; ++i)
{
p_means[k] += *p_src;
p_invstds[k] += (*p_src)*(*p_src);
++p_src;
const float val = *p;
const double delta1 = val - mean;
mean += delta1 / (count + 1);
const double delta2 = val - mean;
M2 += delta1 * delta2;
++count;
++p;
}
}
}
means /= src.num_samples()*num;
invstds /= src.num_samples()*num;
// copy data back to host
invstds.host(); means.host();

p_src = src.host();
// compute variances
running_variances.copy_size(invstds);
auto rvar = running_variances.host();
// This scale makes the running variances unbiased.
const double scale = (src.num_samples()*num)/(src.num_samples()*num-1.0);
for (long k = 0; k < src.k(); ++k)
{
float actual_var = p_invstds[k] - p_means[k]*p_means[k];
const auto actual_var = (count > 1) ? (M2 / count) : 0.0;

if (averaging_factor == 1)
rvar[k] = scale*actual_var;
else
rvar[k] = (1-averaging_factor)*rvar[k] + scale*averaging_factor*actual_var;

p_means[k] = mean;
p_invstds[k] = 1.0f/std::sqrt(actual_var + eps);
}

p_src = src.host();
auto p_dest = dest.host();
for (long n = 0; n < src.num_samples(); ++n)
{
Expand All @@ -1145,7 +1148,6 @@ namespace dlib
}

// now keep track of the running means
running_means.copy_size(means);
if (averaging_factor != 1)
running_means = (1-averaging_factor)*mat(running_means) + averaging_factor*mat(means);
else
Expand Down
4 changes: 2 additions & 2 deletions dlib/test/dnn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -507,7 +507,7 @@ namespace
using namespace dlib::tt;
print_spinner();
resizable_tensor src, gamma, beta, dest, dest2, dest3, means, vars, gradient_input;
src = matrix_cast<float>(gaussian_randm(5,5, 0));
src = matrix_cast<float>(gaussian_randm(5,5, 0) + 10);
gamma = matrix_cast<float>(gaussian_randm(1,5, 1));
beta = matrix_cast<float>(gaussian_randm(1,5, 2));
gradient_input = matrix_cast<float>(gaussian_randm(5,5, 3));
Expand Down Expand Up @@ -593,7 +593,7 @@ namespace
print_spinner();
resizable_tensor src(5,5,4,4), gamma, beta, dest, dest2, dest3, means, vars, gradient_input(5,5,4,4);
tt::tensor_rand rnd;
rnd.fill_gaussian(src);
rnd.fill_gaussian(src,10);
rnd.fill_gaussian(gradient_input);
gamma = matrix_cast<float>(gaussian_randm(1,5, 1));
beta = matrix_cast<float>(gaussian_randm(1,5, 2));
Expand Down
Loading