torch_bitmask

Install

pip install git+https://github.com/mgoin/torch_bitmask

Usage

import torch
from torch_bitmask import NumpyBitmaskTensor as BitmaskTensor
dense_tensor = torch.randn([50, 50])
bitmask_tensor = BitmaskTensor.from_dense(dense_tensor)
decompressed_tensor = bitmask_tensor.to_dense()
assert torch.equal(dense_tensor, decompressed_tensor)

Compression Demo

python demo.py
Generating a tensor of size=[4096, 4096] and precision=torch.float16 with sparsity=0.5
dense_tensor: 32.0000 MB
bitmask_tensor: 18.0175 MB
 values: 15.9863 MB
 bitmask: 2.0000 MB
 row_offsets: 0.0312 MB

Benchmark

python benchmark.py
Generating a tensor of size=[16384, 4096] and precision=torch.float32 with sparsity=0.5
Create Regular Tensor: 1.7972s
Dense memory used: 256.0000 MB
Benchmark NumpyBitmaskTensor:
 compress: 0.1037 sec
 decompress: 0.0848 sec
 memory used: 136.1031 MB
Benchmark TritonBitmaskTensor:
 compress: 0.0697 sec
 decompress: 0.0421 sec
 memory used: 136.1031 MB

Tests

python -m pytest tests 
================================================ test session starts =================================================
platform linux -- Python 3.10.12, pytest-8.0.0, pluggy-1.4.0
rootdir: /home/mgoin/code/bitmask_sparse
plugins: anyio-3.7.1
collected 48 items 
tests/test_compression.py ................................................ [100%]
================================================= 48 passed in 5.61s =================================================