RPC Dataloader

This library implements a variant of the PyTorch Dataloader using remote workers. It allows to distribute workers over remote servers rather than the one running the main script.

To use it, start one or several worker daemons on remote computers. The machines running the data loaders will dispatch requests for items to the workers and await the returned values.

Though similar to torch.rpc, this library uses its own implementation of RPC (Remote Procedure Call) which is simpler (no initialization) and does not conflict with the one from pytorch.

Installation

pip install rpcdataloader

Usage

To use the RPC dataloader, start a few workers either from the command line:

python -m rpcdataloader.launch --host=0.0.0.0 --port=6543

or by calling rpcdataloader.run_worker directly from a python script.

Then instantiate a remote dataset and dataloader:

dataset = rpcdataloader.RPCDataset(
    workers=['node01:6543', 'node02:5432'],
    dataset=torchvision.datasets.ImageFolder,
    root=args.data_path + "/train",
    transform=train_transform,
)

dataloader = rpcdataloader.RPCDataloader(
    dataset
    batch_size=2,
    shuffle=True,
    pin_memory=True)

for minibatch in dataloader:
    ...

Further reading

• API documentation
• ResNet50 training on ImageNet dataset
• Slurm integration using heterogeneous jobs