https://github.com/feifeibear/pstensor

PSTensor provides a way to hack the memory management of tensors in TensorFlow and PyTorch by defining your own C++ Tensor Class.
https://github.com/feifeibear/pstensor

cuda deeplearning machinelearning pytorch tensorflow2

Last synced: 16 days ago
JSON representation

PSTensor provides a way to hack the memory management of tensors in TensorFlow and PyTorch by defining your own C++ Tensor Class.

• Host: GitHub
• URL: https://github.com/feifeibear/pstensor
• Owner: feifeibear
• License: other
• Created: 2021-11-16T09:03:11.000Z (about 3 years ago)
• Default Branch: master
• Last Pushed: 2022-02-10T05:38:46.000Z (almost 3 years ago)
• Last Synced: 2024-12-25T15:27:54.560Z (about 1 month ago)
• Topics: cuda, deeplearning, machinelearning, pytorch, tensorflow2
• Language: C++
• Homepage:
• Size: 31.3 KB
• Stars: 10
• Watchers: 2
• Forks: 0
• Open Issues: 1
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

        
## PSTensor : Customized a Tensor Data Structure Compatible with PyTorch and TensorFlow.

You may need this software in the following cases.

1. Manage memory allocation by yourself. Sometimes, you are irritated by the framework's memory allocation mechanism. They use a complicated caching-based allocator and generate fragments.

2. Unified framework-agnostic memory management operations.

3. Customized Communication Pattern. Using PyTorch, it is impossible to implement GPU P2P communication, since nccl backend only supports collective communication APIs. Now, you can implement it with help of CUDA-level libraries.

## Installation

```

mkdir build && cd build && cmake .. && make

pip install `find . -name "*whl"`

```

## Usage

See [PyTorch Example](./test_torch.py) and [TensorFlow Example](./test_tf.py)  for details.

More features are Working In Progress.