Welcome to run-batched’s documentation!

run-batched is a Python package that allows you to run a torch model on a numpy array in batches. There is one function in this package:

run_batched.run_batched(m, x, bs, pbar=<function <lambda>>, *, device='cuda')

Run m on x in batches of size bs, and concatenate the results.

If x is a dict, treat it as a dict of arrays, each of which is batched together. This requires that the length of each array is the same.

If m returns a dict, treat it as a dict of arrays, each of which is concatenated together. This requires that the length of each array is the same.

Parameters:
  • m – A function which takes a batch of data and returns a batch of data. The inputs and outputs are either tensors or dicts of tensors.

  • x – The data to be batched.

  • bs – The batch size.

  • pbar – A function which takes an iterable and returns an iterable. This is used to wrap the iterable over the batches of data to produce a pbar.

Returns:

The result of running m on x in batches of size bs.