batch-dev 0.0.3

Description:

batchdev 0.0.3

Batch
Generic python module for handling dictionary-based batch data.





Purpose
Are you working with data of similar modalities, and often have to apply the same function to multiple elements?
Are you using something similar to this:
batch = {
"image_a": image_a,
"image_b": image_b,
"image_c": image_c
}

# Move to another device
for key in batch:
batch[key] = batch[key].to(device)

# Transform
for key in batch:
batch[key] = batch[key] * 2 + 1

# Combine
for key in batch:
batch[key] = batch[key] + batch_2[key]

# Process
for key in batch:
batch[key] = batch[key].max()

If the answer is yes, then this module is for you!
Our Batch package is a generic wrapper for dictionary-based batch data.
It provides a simple way to apply the same function or operator to the whole batch.
The module is completely device and container independent, you can use it with PyTorch, NumPy or any other libraries.
batch = Batch(
image_a=image_a,
image_b=image_b,
image_c=image_c)

# Move to another device
batch = batch.to(device)

# Transform
batch = batch * 2 + 1

# Combine
batch = batch + batch_2

# Process
batch = batch.max()

Installation
pip install batch

Usage
The example below demonstrates a few basic use-cases using NumPy. Similarly, PyTorch or other containers can also be used.
Import
from batch import Batch

Instantiation
Direct
# Create a batch directly
batch = Batch(
image_a=np.random.rand(256, 256, 3),
image_b=np.random.rand(256, 256, 3),
image_c=np.random.rand(256, 256, 3))

From dictionary
batch = {
"image_a": np.random.rand(256, 256, 3),
"image_b": np.random.rand(256, 256, 3),
"image_c": np.random.rand(256, 256, 3)}
# Create a batch from a dictionary
batch = Batch.from_dict(batch)

From tensor
image = np.random.rand(256, 256, 9)

# Create a batch from a tensor by splitting the tensor along one dimension and store the splits
data_splits = {
"image_a": 1,
"image_b": 1,
"image_c": 1}
dim = 2
batch = Batch.from_tensor(batch, data_splits, dim=dim)

Indexing
A Batch is a string-keyed dictionary, with potentially mapping or iterable values.
String index
When a string index is given, then it is always interpreted as a key.
Single key
Querying a single key returns the value associated:
image_a = batch["image_a"]

You can even index deeper using . as separator:
batch_2 = Batch(input=batch)
image_a = batch_2["input.image_a"]

Multiple keys
Querying multiple keys (tuple or list) return a new batch with the selected keys:
batch_out = batch["image_a", "image_b"]

Wildcard query
Wildcard query is also supported and returns a new batch with the matching keys:
batch_out = batch.query_wildcard("image_*")

Integer index
When an integer index is given, then it is always interpreted as an index to the elements and returns a new batch with the indexed elements:
batch_out = batch[:,:,0]

Processing a batch
Operators
You can use the followingunary, binary and reverse operators:
# Unary operators
"__not__", "__abs__", "__index__", "__inv__", "__invert__", "__neg__", "__pos__",

# Binary operators
"__add__", "__and__", "__concat__", "__floordiv__", "__lshift__", "__mod__", "__mul__",
"__or__", "__pow__", "__rshift__", "__sub__", "__truediv__", "__xor__", "__eq__",

# Reverse operators
"__radd__", "__rand__", "__rmul__",
"__ror__", "__rsub__", "__rxor__",

# In-place operators
"__iadd__", "__iand__", "__iconcat__", "__ifloordiv__", "__ilshift__", "__imod__", "__imul__",
"__ior__", "__ipow__", "__irshift__", "__isub__", "__itruediv__", "__ixor__"

Example:
# Use operators
batch_out = batch + batch_2 * 2

Member functions
You can use any member functions of the underlying container, for example:
# Use member functions
batch_out = batch.mean(axis=2)

Map
You can easily apply a function to the whole batch:
batch = batch.map(list) # Converts all elements to list
batch = batch.map(np.stack, axis=0) # Concatenates all elements to a single tensor

Map keys
You can also apply a function to the keys:
batch = batch.map_keys(lambda x: f"{x}_2") # Add suffix

Limitations
A few limitations to consider when using this module:

Use only string keys for the batch.
Don't use keys starting with underscore (_).
Slice indexing is not implemented yet.
Generic iterable indexing is not implemented, only tuple and list.
Code documentation is in progress.
Some features are not yet documented here, please refer to the code directly.

If you have any ideas or requests, feel free to open an issue or a pull request.