tensorblob 0.1.3

Lightweight, dynamic-sized, memory-mapped tensor storage with file-like APIs.

tensorblob

A lightweight, dynamic-sized, memory-mapped tensor storage with file-like APIs, while also supporting integer indexing and slicing, built with MemoryMappedTensor from tensordict.

Features

• 🔗 Memory-mapped storage: Efficient storage of large collections of same-shaped tensors
• 💾 File-like APIs: Read, write, and seek like a file, while also supporting integer indexing and slicing
• ⚡ Dynamic-sized: No need to specify the total number of tensors upfront
• 🔄 Extend and truncate: Extend the blob with another blob or truncate the blob to a specific position
• 🚀 LRU cache: Automatic management of memory-mapped blocks for scalability with large blobs

Installation

From PyPI:

pip install tensorblob

If you are interested in the experimental (i.e., unstable and undertested) version, you can install it from GitHub:

pip install git+https://github.com/Guest400123064/tensorblob.git

Core Use Cases

Quick Start

The example below shows how to create a new storage for a collection of randomly generated fake embeddings, and how to access them by index. Since the storage is memory-mapped, no need to read all tensors into memory; just access them by index.

import torch
from tensorblob import TensorBlob

# Create a new storage for a collection of randomly generated fake embeddings;
# need to specify the data type and shape of each tensor for creation
with TensorBlob.open("embeddings.blob", "w", dtype="float32", shape=768) as blob:
    blob.write(torch.randn(100_000, 768))
    print(f"Wrote {len(blob)} embeddings")

# No need to specify the configurations again after creation
with TensorBlob.open("embeddings.blob", "r") as blob:
    e1 = blob[42]
    e2 = blob[-1:16384:-12345]
    print(f"Similarity: {torch.cosine_similarity(e1, e2)}")

Processing Large Datasets

Store and preprocess datasets larger than RAM using memory mapping can be useful to accelerate the training process by reducing the time spent on data loading and transformation.

with TensorBlob.open("data/images.blob", "w", dtype="float32", shape=(3, 224, 224)) as blob:
    for image_batch in data_loader:
        blob.write(preprocess(image_batch))

with TensorBlob.open("data/images.blob", "r") as blob:
    for image in blob:
        result = model(image)

Incremental Data Collection

Append new data to existing blobs can be useful with streaming data collection.

with TensorBlob.open("positions.blob", "w", dtype="float32", shape=3) as blob:
    blob.write(initial_position)

# Later: append more data by opening the blob in append mode
with TensorBlob.open("positions.blob", "a") as blob:
    for pos in trajectory_queue.get():
        blob.write(pos)
    print(f"Total trajectory recorded: {len(blob)}")

Random Access and Updates with File-Like APIs

Read and modify specific tensors starting from a specific position.

import io

with TensorBlob.open("data/features.blob", "r+") as blob:
    blob.seek(1000)
    print(f"Current position: {blob.tell()}")

    batch = blob.read(size=100)
    print(f"Read {batch.shape} tensors")

    # Update specific positions, whence is also supported
    blob.seek(-500, whence=io.SEEK_END)
    blob.write(updated_features)
    
    # Append new data
    blob.seek(len(blob))
    blob.write(additional_features)

Extend and Truncate

Extend the blob with another blob or truncate the blob to a specific position. Extension could be useful if we want to merge two blobs into one, e.g., results from two different processes. Note that extension operation does not delete the original data.

with TensorBlob.open("data/features.blob", "a") as blob:
    blob.extend(other_blob)

# Extension without maintaining the order is faster
with TensorBlob.open("data/features.blob", "r+") as blob:
    blob.extend(other_blob, maintain_order=False)

with TensorBlob.open("data/features.blob", "r+") as blob:
    blob.truncate(1000)
    print(f"Truncated to {len(blob)} tensors")

Performance and Scalability

Memory Management

TensorBlob uses an LRU (Least Recently Used) cache to manage memory-mapped blocks efficiently. This allows you to work with blobs containing millions of tensors without loading everything into memory.

Default behavior:

• Automatically caches up to ~4,000 blocks (1/16 of system's VMA limit)
• Blocks loaded on-demand when accessed
• Least recently used blocks automatically evicted when cache is full

For large-scale workloads:

# Increase cache for better random access performance
with TensorBlob.open("large.blob", "r", max_cached_blocks=10_000) as blob:
    for idx in random_indices:
        tensor = blob[idx]  # Cached blocks reused efficiently

# Decrease cache for memory-constrained environments
with TensorBlob.open("data.blob", "r", max_cached_blocks=100) as blob:
    for tensor in blob:  # Sequential access works fine with small cache
        process(tensor)

Performance tips:

• Sequential access patterns work well with any cache size
• Random access benefits from larger cache sizes
• Each cached block consumes ~200 bytes of kernel memory (VMA overhead)
• System limit: typically ~65,000 memory-mapped regions per process
• To avoid frequent cache evictions, one can also increase the block size to reduce the total number of blocks

Contributing

Contributions welcome! Please submit a Pull Request.

License

Apache License 2.0 - see LICENSE file for details.