Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
FAST-RNNT
Commits
89beacff
"git@developer.sourcefind.cn:OpenDAS/ollama.git" did not exist on "9502e5661f97287064e264091b3bbc88646ac1fb"
Commit
89beacff
authored
Jul 02, 2021
by
Daniel Povey
Browse files
Remove debug statements
parent
b3c1340d
Changes
2
Hide whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
50 additions
and
44 deletions
+50
-44
README.md
README.md
+46
-44
torch_integrated_conv/integrated_conv_cuda_kernel.cu
torch_integrated_conv/integrated_conv_cuda_kernel.cu
+4
-0
No files found.
README.md
View file @
89beacff
# TODO: update this README!
# Fast Discounted Cumulative Sums in PyTorch
# Fast Discounted Cumulative Sums in PyTorch
[

](https://pypi.org/project/torch-discounted-cumsum/)
[

](https://pypi.org/project/torch-discounted-cumsum/)
...
@@ -9,15 +11,15 @@
...
@@ -9,15 +11,15 @@
<img
src=
"doc/img/logo_small.png"
align=
"left"
width=
"33%"
>
<img
src=
"doc/img/logo_small.png"
align=
"left"
width=
"33%"
>
This repository implements an efficient parallel algorithm for the computation of discounted cumulative sums
This repository implements an efficient parallel algorithm for the computation of discounted cumulative sums
and a Python package with differentiable bindings to PyTorch. The discounted
`cumsum`
operation is frequently seen in
and a Python package with differentiable bindings to PyTorch. The discounted
`cumsum`
operation is frequently seen in
data science domains concerned with time series, including Reinforcement Learning (RL).
data science domains concerned with time series, including Reinforcement Learning (RL).
The traditional sequential algorithm performs the computation of the output elements in a loop. For an input of size
The traditional sequential algorithm performs the computation of the output elements in a loop. For an input of size
`N`
, it requires
`O(N)`
operations and takes
`O(N)`
time steps to complete.
`N`
, it requires
`O(N)`
operations and takes
`O(N)`
time steps to complete.
The proposed parallel algorithm requires a total of
`O(N log N)`
operations, but takes only
`O(log N)`
time steps, which is a
The proposed parallel algorithm requires a total of
`O(N log N)`
operations, but takes only
`O(log N)`
time steps, which is a
considerable trade-off in many applications involving large inputs.
considerable trade-off in many applications involving large inputs.
Features of the parallel algorithm:
Features of the parallel algorithm:
-
Speed logarithmic in the input size
-
Speed logarithmic in the input size
...
@@ -74,7 +76,7 @@ K = 2
...
@@ -74,7 +76,7 @@ K = 2
gamma
=
0.99
gamma
=
0.99
x
=
torch
.
ones
(
1
,
N
).
cuda
()
x
=
torch
.
ones
(
1
,
N
).
cuda
()
y_N
=
discounted_cumsum_right
(
x
,
gamma
)
y_N
=
discounted_cumsum_right
(
x
,
gamma
)
y_K
=
y_N
-
(
gamma
**
K
)
*
torch
.
cat
((
y_N
[:,
K
:],
torch
.
zeros
(
1
,
K
).
cuda
()),
dim
=
1
)
y_K
=
y_N
-
(
gamma
**
K
)
*
torch
.
cat
((
y_N
[:,
K
:],
torch
.
zeros
(
1
,
K
).
cuda
()),
dim
=
1
)
print
(
y_K
)
print
(
y_K
)
```
```
...
@@ -88,68 +90,68 @@ tensor([[1.9900, 1.9900, 1.9900, 1.9900, 1.9900, 1.9900, 1.9900, 1.0000]],
...
@@ -88,68 +90,68 @@ tensor([[1.9900, 1.9900, 1.9900, 1.9900, 1.9900, 1.9900, 1.9900, 1.0000]],
## Parallel Algorithm
## Parallel Algorithm
For the sake of simplicity, the algorithm is explained for
`N=16`
.
For the sake of simplicity, the algorithm is explained for
`N=16`
.
The processing is performed in-place in the input vector in
`log2 N`
stages. Each stage updates
`N / 2`
positions in parallel
The processing is performed in-place in the input vector in
`log2 N`
stages. Each stage updates
`N / 2`
positions in parallel
(that is, in a single time step, provided unrestricted parallelism). A stage is characterized by the size of the group of
(that is, in a single time step, provided unrestricted parallelism). A stage is characterized by the size of the group of
sequential elements being updated, which is computed as
`2 ^ (stage - 1)`
.
sequential elements being updated, which is computed as
`2 ^ (stage - 1)`
.
The group stride is always twice larger than the group size. The elements updated during the stage are highlighted with
The group stride is always twice larger than the group size. The elements updated during the stage are highlighted with
the respective stage color in the figure below. Here input elements are denoted with their position id in hex, and the
the respective stage color in the figure below. Here input elements are denoted with their position id in hex, and the
elements tagged with two symbols indicate the range over which the discounted partial sum is computed upon stage completion.
elements tagged with two symbols indicate the range over which the discounted partial sum is computed upon stage completion.
Each element update includes an in-place addition of a discounted element, which follows the last
Each element update includes an in-place addition of a discounted element, which follows the last
updated element in the group. The discount factor is computed as gamma raised to the power of the distance between the
updated element in the group. The discount factor is computed as gamma raised to the power of the distance between the
updated and the discounted elements. In the figure below, this operation is denoted with tilted arrows with a greek
updated and the discounted elements. In the figure below, this operation is denoted with tilted arrows with a greek
gamma tag. After the last stage completes, the output is written in place of the input.
gamma tag. After the last stage completes, the output is written in place of the input.
<p
align=
"center"
>
<p
align=
"center"
>
<img
src=
"doc/img/algorithm.png"
>
<img
src=
"doc/img/algorithm.png"
>
</p>
</p>
In the CUDA implementation,
`N / 2`
CUDA threads are allocated during each stage to update the respective elements. The
In the CUDA implementation,
`N / 2`
CUDA threads are allocated during each stage to update the respective elements. The
strict separation of updates into stages via separate kernel invocations guarantees stage-level synchronization and
strict separation of updates into stages via separate kernel invocations guarantees stage-level synchronization and
global consistency of updates.
global consistency of updates.
The gradients wrt input can be obtained from the gradients wrt output by simply taking
`cumsum`
operation with the
The gradients wrt input can be obtained from the gradients wrt output by simply taking
`cumsum`
operation with the
reversed direction of summation.
reversed direction of summation.
## Numerical Precision
## Numerical Precision
The parallel algorithm produces a more numerically-stable output than the sequential algorithm using the same scalar
The parallel algorithm produces a more numerically-stable output than the sequential algorithm using the same scalar
data type.
data type.
The comparison is performed between 3 runs with identical inputs (
[
code
](
tests/test.py#L116
)
). The first run casts inputs to
The comparison is performed between 3 runs with identical inputs (
[
code
](
tests/test.py#L116
)
). The first run casts inputs to
double precision and obtains the output reference using the sequential algorithm. Next, we run both sequential and
double precision and obtains the output reference using the sequential algorithm. Next, we run both sequential and
parallel algorithms with the same inputs cast to single precision and compare the results to the reference. The
parallel algorithms with the same inputs cast to single precision and compare the results to the reference. The
comparison is performed using the
`L_inf`
norm, which is just the maximum of per-element discrepancies.
comparison is performed using the
`L_inf`
norm, which is just the maximum of per-element discrepancies.
With 10000-element non-zero-centered input (such as all elements are 1.0), the errors of the algorithms are 2.8e-4
With 10000-element non-zero-centered input (such as all elements are 1.0), the errors of the algorithms are 2.8e-4
(sequential) and 9.9e-5 (parallel). With zero-centered inputs (such as standard gaussian noise), the errors are
(sequential) and 9.9e-5 (parallel). With zero-centered inputs (such as standard gaussian noise), the errors are
1.
8e-5 (sequential) and 1.5e-5 (parallel).
1.
8e-5 (sequential) and 1.5e-5 (parallel).
## Speed-up
## Speed-up
We tested 3 implementations of the algorithm with the same 100000-element input (
[
code
](
tests/test.py#L154
)
):
We tested 3 implementations of the algorithm with the same 100000-element input (
[
code
](
tests/test.py#L154
)
):
1.
Sequential in PyTorch on CPU (as in
1.
Sequential in PyTorch on CPU (as in
[
REINFORCE
](
https://github.com/pytorch/examples/blob/87d9a1e/reinforcement_learning/reinforce.py#L66-L68
)
) (Intel Xeon CPU, DGX-1)
[
REINFORCE
](
https://github.com/pytorch/examples/blob/87d9a1e/reinforcement_learning/reinforce.py#L66-L68
)
) (Intel Xeon CPU, DGX-1)
2.
Sequential in C++ on CPU (Intel Xeon CPU, DGX-1)
2.
Sequential in C++ on CPU (Intel Xeon CPU, DGX-1)
3.
Parallel in CUDA (NVIDIA P-100, DGX-1)
3.
Parallel in CUDA (NVIDIA P-100, DGX-1)
The observed speed-ups are as follows:
The observed speed-ups are as follows:
-
PyTorch to C++: 387 times
-
PyTorch to C++: 387 times
-
PyTorch to CUDA: 36573 times
-
PyTorch to CUDA: 36573 times
-
C++ to CUDA: 94 times
-
C++ to CUDA: 94 times
## Ops-Space-Time Complexity
## Ops-Space-Time Complexity
Assumptions:
Assumptions:
-
A fused operation of raising
`gamma`
to a power, multiplying the result by
`x`
, and adding
`y`
is counted as a
-
A fused operation of raising
`gamma`
to a power, multiplying the result by
`x`
, and adding
`y`
is counted as a
single fused operation;
single fused operation;
-
`N`
is a power of two. When it isn't, the parallel algorithm's complexity is the same as with N equal to the next
-
`N`
is a power of two. When it isn't, the parallel algorithm's complexity is the same as with N equal to the next
power of two.
power of two.
Under these assumptions, the sequential algorithm takes
`N`
operations and
`N`
time steps to complete.
Under these assumptions, the sequential algorithm takes
`N`
operations and
`N`
time steps to complete.
The parallel algorithm takes
`0.5 * N * log2 N`
operations and can be completed in
`log2 N`
time steps
The parallel algorithm takes
`0.5 * N * log2 N`
operations and can be completed in
`log2 N`
time steps
if the parallelism is unrestricted.
if the parallelism is unrestricted.
Both algorithms can be performed in-place; hence their space complexity is
`O(1)`
.
Both algorithms can be performed in-place; hence their space complexity is
`O(1)`
.
...
@@ -157,20 +159,20 @@ Both algorithms can be performed in-place; hence their space complexity is `O(1)
...
@@ -157,20 +159,20 @@ Both algorithms can be performed in-place; hence their space complexity is `O(1)
#### PyTorch
#### PyTorch
As of the time of writing, PyTorch does not provide discounted
`cumsum`
functionality via the API. PyTorch RL code
As of the time of writing, PyTorch does not provide discounted
`cumsum`
functionality via the API. PyTorch RL code
samples (e.g.,
[
REINFORCE
](
https://github.com/pytorch/examples/blob/87d9a1e/reinforcement_learning/reinforce.py#L66-L68
)
)
samples (e.g.,
[
REINFORCE
](
https://github.com/pytorch/examples/blob/87d9a1e/reinforcement_learning/reinforce.py#L66-L68
)
)
suggest computing returns in a loop over reward items. Since most RL algorithms do not require differentiating through
suggest computing returns in a loop over reward items. Since most RL algorithms do not require differentiating through
returns, many code samples resort to using SciPy function listed below.
returns, many code samples resort to using SciPy function listed below.
#### TensorFlow
#### TensorFlow
TensorFlow API provides
`tf.scan`
API, which can be supplied with an appropriate lambda function implementing the
TensorFlow API provides
`tf.scan`
API, which can be supplied with an appropriate lambda function implementing the
formula above. Under the hood, however,
`tf.scan`
implement the traditional sequential algorithm.
formula above. Under the hood, however,
`tf.scan`
implement the traditional sequential algorithm.
#### SciPy
#### SciPy
SciPy provides a
`scipy.signal.lfilter`
function for computing IIR filter response using the sequential algorithm, which
SciPy provides a
`scipy.signal.lfilter`
function for computing IIR filter response using the sequential algorithm, which
can be used for the task at hand, as suggested in this
[
StackOverflow
](
https://stackoverflow.com/a/47971187/411907
)
can be used for the task at hand, as suggested in this
[
StackOverflow
](
https://stackoverflow.com/a/47971187/411907
)
response.
response.
## Citation
## Citation
...
...
torch_integrated_conv/integrated_conv_cuda_kernel.cu
View file @
89beacff
...
@@ -665,6 +665,7 @@ torch::Tensor integrated_conv_cuda(torch::Tensor input,
...
@@ -665,6 +665,7 @@ torch::Tensor integrated_conv_cuda(torch::Tensor input,
assert
(
num_blocks_patch
<=
num_patches
&&
num_blocks_batch
<=
N
);
assert
(
num_blocks_patch
<=
num_patches
&&
num_blocks_batch
<=
N
);
#if 0
std::cout << "N,C,H,W=" << N << "," << C << "," << H << "," << W
std::cout << "N,C,H,W=" << N << "," << C << "," << H << "," << W
<< "; kW,kH=" << kW << "," << kH
<< "; kW,kH=" << kW << "," << kH
<< "; patchH,patchW=" << patchH << ","
<< "; patchH,patchW=" << patchH << ","
...
@@ -674,6 +675,7 @@ torch::Tensor integrated_conv_cuda(torch::Tensor input,
...
@@ -674,6 +675,7 @@ torch::Tensor integrated_conv_cuda(torch::Tensor input,
<< ", threads_per_opixel=" << threads_per_opixel
<< ", threads_per_opixel=" << threads_per_opixel
<< ", threads_per_block=" << threads_per_block
<< ", threads_per_block=" << threads_per_block
<< std::endl;
<< std::endl;
#endif
dim3
gridDim
(
C
,
num_blocks_patch
,
num_blocks_batch
);
dim3
gridDim
(
C
,
num_blocks_patch
,
num_blocks_batch
);
// blockDim is scalar, just threads_per_block.
// blockDim is scalar, just threads_per_block.
...
@@ -805,6 +807,7 @@ std::vector<torch::Tensor> integrated_conv_backward_cuda(torch::Tensor input,
...
@@ -805,6 +807,7 @@ std::vector<torch::Tensor> integrated_conv_backward_cuda(torch::Tensor input,
assert
(
patchH
*
patchW
*
threads_per_pixel
<=
threads_per_block
);
assert
(
patchH
*
patchW
*
threads_per_pixel
<=
threads_per_block
);
assert
(
kH
*
kW
*
threads_per_kernel_pos
<=
threads_per_block
);
assert
(
kH
*
kW
*
threads_per_kernel_pos
<=
threads_per_block
);
#if 0
std::cout << "[backward:] N,C,H,W=" << N << "," << C << "," << H << "," << W
std::cout << "[backward:] N,C,H,W=" << N << "," << C << "," << H << "," << W
<< "; kW,kH=" << kW << "," << kH
<< "; kW,kH=" << kW << "," << kH
<< "; patchH,patchW=" << patchH << ","
<< "; patchH,patchW=" << patchH << ","
...
@@ -816,6 +819,7 @@ std::vector<torch::Tensor> integrated_conv_backward_cuda(torch::Tensor input,
...
@@ -816,6 +819,7 @@ std::vector<torch::Tensor> integrated_conv_backward_cuda(torch::Tensor input,
<< ", threads_per_block=" << threads_per_block
<< ", threads_per_block=" << threads_per_block
<< ", buffer_numel=" << buffer_numel
<< ", buffer_numel=" << buffer_numel
<< std::endl;
<< std::endl;
#endif
int
num_blocks
=
num_blocks_patch
*
num_blocks_batch
;
int
num_blocks
=
num_blocks_patch
*
num_blocks_batch
;
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment