Skip to content

GitLab

Commit 7ce3c947 authored by Daniel Povey's avatar Daniel Povey
Browse files

Revise CPU backward code, simpler now.

parent fa53fa33
Show whitespace changes
Inline Side-by-side
Showing with 34 additions and 52 deletions
torch_learned_nonlin/learned_nonlin_cpu.cpp
View file @ 7ce3c947
......@@ -44,8 +44,6 @@ torch::Tensor learned_nonlin_cpu(torch::Tensor input,
sum_negative -= neg_scaled_param;
y_vals_a[c][K - i - 1] = sum_negative + neg_scaled_param * (i + 1);
}
//scalar_t neg_scaled_param = params_a[c][1] * scale;
//y_vals_a[c][0] = sum_negative + neg_scaled_param * K;
}
auto input_a = input.accessor<scalar_t, 3>(),
......@@ -68,10 +66,7 @@ torch::Tensor learned_nonlin_cpu(torch::Tensor input,
// C++ rounds toward zero.
int n = (int) x;
// OK, at this point, 0 <= min < 2*K.
scalar_t y = input * params_a[c][n + 1] + y_vals_a[c][n];
/* printf("x = %f, y = %f, n = %d; y = (%f - %d) * %f+ %f\n", x, y, n,
x, n, params_a[c][n + 1], y_vals_a[c][n - 1]); */
output_a[b][c][t] = y;
output_a[b][c][t] = input * params_a[c][n + 1] + y_vals_a[c][n];
}
}
}}));
......@@ -122,17 +117,14 @@ std::vector<torch::Tensor> learned_nonlin_backward_cpu(torch::Tensor input,
sum_positive = 0.0,
scale = exp(params_a[c][0]);
for (int i = 0; i < K; i++) {
y_vals_a[c][K + i] = sum_positive;
y_vals_a[c][K - i] = sum_negative;
sum_positive += params_a[c][1 + K + i] * scale;
sum_negative -= params_a[c][K - i] * scale;
scalar_t pos_scaled_param = params_a[c][1 + K + i] * scale,
neg_scaled_param = params_a[c][K - i] * scale;
y_vals_a[c][K + i] = sum_positive - pos_scaled_param * i;
sum_positive += pos_scaled_param;
sum_negative -= neg_scaled_param;
y_vals_a[c][K - i - 1] = sum_negative + neg_scaled_param * (i + 1);
}
// the reference point for the lowest, half-infinite interval (the one
// starting at x=-(K-1) is x=-K; this is arbitrary but makes the
// computation more regular.
y_vals_a[c][0] = sum_negative;
}
auto input_a = input.accessor<scalar_t, 3>(),
output_grad_a = output_grad.accessor<scalar_t, 3>(),
input_grad_a = input_grad.accessor<scalar_t, 3>();
......@@ -144,57 +136,47 @@ std::vector<torch::Tensor> learned_nonlin_backward_cpu(torch::Tensor input,
inv_scale_grad = 0.0,
scale_grad = 0.0;
for (int t = 0; t < T; t++) {
// `x` is the scaled input x plus an offset so that -K maps to 0.
// Note: the discontinuities in our function are at -(K-1) ... +(K+1),
// so in a sense -K and +K are not special, but we include those
// extra values as an easy way to handle the semi-infinite regions
// that are < -(K-1) and > (K-1)
scalar_t input = input_a[b][c][t],
x = input * inv_scale + K,
y_grad = output_grad_a[b][c][t],
x_trunc = x;
if (x_trunc < 0) x_trunc = 0;
else if (x_trunc >= N) x_trunc = N - 1;
output_grad = output_grad_a[b][c][t];
if (x < 0) x = 0;
else if (x >= N) x = N - 1;
// C++ rounds toward zero.
int n = (int) x_trunc;
int n = (int) x;
// OK, at this point, 0 <= n < 2*K.
// backprop for:
// scalar_t x_residual_scaled = (x - (scalar_t)n) * scale
// scalar_t y = x_residual_scaled * params_a[c][n + 1] + y_vals_a[c][n];
scalar_t x_residual_scaled = (x - n) * scale,
x_residual_scaled_grad = y_grad * params_a[c][n + 1],
x_grad = x_residual_scaled_grad * scale;
scale_grad += x_residual_scaled_grad * (x - (scalar_t)n);
params_grad_a[c][n + 1] += y_grad * x_residual_scaled;
y_vals_grad_a[c][n] += y_grad;
// backprop for: x = input * inv_scale + K,
inv_scale_grad += x_grad * input;
input_grad_a[b][c][t] = x_grad * inv_scale;
// output_a[b][c][t] = input * params_a[c][n + 1] + y_vals_a[c][n];
params_grad_a[c][n + 1] += output_grad * input;
y_vals_grad_a[c][n] += output_grad;
input_grad_a[b][c][t] = output_grad * params_a[c][n + 1];
}
// Do the backprop for:
// scale = exp(params_a[c][0]);
// inv_scale = exp(-params_a[c][0]);
params_grad_a[c][0] += (scale * scale_grad - inv_scale * inv_scale_grad);
}
}
// Now do the backprop for the loop above where we set y_vals_a.
for (int c = 0; c < C; c++) {
scalar_t scale = exp(params_a[c][0]),
scale_grad = 0.0,
sum_negative_grad = y_vals_grad_a[c][0], // backprop for: y_vals_a[c][0] = sum_negative
sum_negative_grad = 0.0,
sum_positive_grad = 0.0;
for (int i = K - 1; i >= 0; i--) {
// backprop for: sum_negative -= params_a[c][K - i] * scale;
params_grad_a[c][K - i] -= sum_negative_grad * scale;
// backprop for: sum_positive += params_a[c][1 + K + i] * scale;
params_grad_a[c][1 + K + i] += sum_positive_grad * scale;
// .. and the contributions to scale_grad for the 2 expressions above..
scale_grad += (sum_positive_grad * params_a[c][1 + K + i] -
sum_negative_grad * params_a[c][K - i]);
// backprop for: y_vals_a[c][K - i] = sum_negative
sum_negative_grad += y_vals_grad_a[c][K - i];
// backprop for: y_vals_a[c][K + i] = sum_positive
sum_positive_grad += y_vals_grad_a[c][K + i];
// Backprop for: y_vals_a[c][K - i - 1] = sum_negative + neg_scaled_param * (i + 1):
scalar_t y_grad_neg = y_vals_grad_a[c][K - i - 1];
sum_negative_grad += y_grad_neg;
scalar_t neg_scaled_param_grad = y_grad_neg * (i + 1);
// Backprop for: sum_negative -= neg_scaled_param;
neg_scaled_param_grad -= sum_negative_grad;
// Backprop for: sum_positive += pos_scaled_param;
scalar_t pos_scaled_param_grad = sum_positive_grad;
// Backprop for: y_vals_a[c][K + i] = sum_positive - pos_scaled_param * i;
scalar_t y_grad_pos = y_vals_grad_a[c][K + i];
pos_scaled_param_grad -= i * y_grad_pos;
sum_positive_grad += y_grad_pos;
// Backprop for: pos_scaled_param = params_a[c][1 + K + i] * scale,
// and: neg_scaled_param = params_a[c][K - i] * scale;
params_grad_a[c][1 + K + i] += pos_scaled_param_grad * scale;
params_grad_a[c][K - i] += neg_scaled_param_grad * scale;
scale_grad += (pos_scaled_param_grad * params_a[c][1 + K + i] +
neg_scaled_param_grad * params_a[c][K - i]);
}
// Backprop for: scale = exp(params_a[c][0]),
params_grad_a[c][0] += scale * scale_grad;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment