Commit 371e2657 authored by Daniel Povey's avatar Daniel Povey
Browse files

Start changing some names

parent 9c56e510
...@@ -4,7 +4,7 @@ excluding any dependencies. ...@@ -4,7 +4,7 @@ excluding any dependencies.
BSD 3-Clause License BSD 3-Clause License
Copyright (c) 2021, Anton Obukhov Copyright (c) 2021, Xiaomi Corporation (written by: Daniel Povey)
All rights reserved. All rights reserved.
Redistribution and use in source and binary forms, with or without Redistribution and use in source and binary forms, with or without
......
include requirements.txt include requirements.txt
include pyproject.toml include pyproject.toml
include LICENSE* include LICENSE*
recursive-include torch_discounted_cumsum * recursive-include torch_mutual_information *
recursive-include doc/img * recursive-include doc/img *
recursive-include tests * recursive-include tests *
global-exclude *.pyc global-exclude *.pyc
\ No newline at end of file
...@@ -9,48 +9,31 @@ with open('requirements.txt') as f: ...@@ -9,48 +9,31 @@ with open('requirements.txt') as f:
long_description = """ long_description = """
This package implements an efficient parallel algorithm for the computation of discounted cumulative sums This package implements an efficient parallel algorithm for the computation of
with differentiable bindings to PyTorch. The `cumsum` operation is frequently seen in data science mutual information between sequences with differentiable bindings to PyTorch.
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
`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, which is a
considerable trade-off in many applications involving large inputs.
Features of the parallel algorithm:
- Speed logarithmic in the input size
- Better numerical precision than sequential algorithms
Features of the package:
- CPU: sequential algorithm in C++
- GPU: parallel algorithm in CUDA
- Gradients computation wrt input
- Both left and right directions of summation supported
- PyTorch bindings
Find more details and the most up-to-date information on the project webpage: Find more details and the most up-to-date information on the project webpage:
https://www.github.com/toshas/torch-discounted-cumsum [TODO]
""" """
def configure_extensions(): def configure_extensions():
out = [ out = [
CppExtension( CppExtension(
'torch_learned_nonlin_cpu', 'torch_mutual_information_cpu',
[ [
os.path.join('torch_learned_nonlin', 'learned_nonlin_cpu.cpp'), os.path.join('torch_mutual_information', 'mutual_information_cpu.cpp'),
], ],
) )
] ]
try: try:
out.append( out.append(
CUDAExtension( CUDAExtension(
'torch_learned_nonlin_cuda', 'torch_mutual_information_cuda',
[ [
os.path.join('torch_learned_nonlin', 'learned_nonlin_cuda.cpp'), os.path.join('torch_mutual_information', 'mutual_information_cuda.cpp'),
os.path.join('torch_learned_nonlin', 'learned_nonlin_cuda_kernel.cu'), os.path.join('torch_mutual_information', 'mutual_information_cuda_kernel.cu'),
], ],
) )
) )
...@@ -60,9 +43,9 @@ def configure_extensions(): ...@@ -60,9 +43,9 @@ def configure_extensions():
setup( setup(
name='torch_learned_nonlin', name='torch_mutual_information',
version='1.0.2', version='1.0.2',
description='Fast discounted cumulative sums in PyTorch', description='Mutual information between sequences of vectors',
long_description=long_description, long_description=long_description,
long_description_content_type='text/markdown', long_description_content_type='text/markdown',
install_requires=requirements, install_requires=requirements,
...@@ -70,13 +53,11 @@ setup( ...@@ -70,13 +53,11 @@ setup(
packages=find_packages(), packages=find_packages(),
author='Dan Povey', author='Dan Povey',
license='BSD', license='BSD',
url='https://www.github.com/toshas/torch-discounted-cumsum',
ext_modules=configure_extensions(), ext_modules=configure_extensions(),
cmdclass={ cmdclass={
'build_ext': BuildExtension 'build_ext': BuildExtension
}, },
keywords=[ keywords=[
'pytorch', 'discounted', 'cumsum', 'cumulative', 'sum', 'scan', 'differentiable', 'pytorch', 'sequence', 'mutual', 'information'
'reinforcement', 'learning', 'rewards', 'time', 'series'
], ],
) )
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