PyTorch: Loading Data from AIStore

Note: The torchdata.datapipes module has been deprecated and removed in recent versions of torchdata. Some information in this blog post may be outdated.

Listing and loading data from AIS buckets (buckets that are not 3rd party backend-based) and remote cloud buckets (3rd party backend-based cloud buckets) using AISFileLister and AISFileLoader.

AIStore (AIS for short) fully supports Amazon S3, Google Cloud, and Microsoft Azure backends, providing a unified namespace across multiple connected backends and/or other AIS clusters, and more.

In the following example, we use the Caltech-256 Object Category Dataset containing 256 object categories and a total of 30607 images stored on an AIS bucket and the Microsoft COCO Dataset which has 330K images with over 200K labels of more than 1.5 million object instances across 80 object categories stored on Google Cloud.

1
# Imports
2
import os
3
from IPython.display import Image
4
5
from torchdata.datapipes.iter import AISFileLister, AISFileLoader, Mapper

Running the AIStore Cluster

Getting started with AIS will take only a few minutes (prerequisites boil down to having a Linux with a disk). See the available deployment options here.

To keep this example simple, we will be running a minimal container-based deployment of AIStore.

1
# Running the AIStore cluster in a container on port 51080
2
# Note: The docker volume must have enough space to load the dataset
3
4
cd deploy/prod/docker/compose
5
6
# Set up environment variables
7
export GCP_CREDENTIALS_PATH=~/.gcp/gcp.json
8
export AWS_CREDENTIALS_PATH=~/.aws 
9
export AWS_REGION="us-east-1"
10
11
# Bring up the cluster
12
make up

To create and put objects (dataset) in the bucket, I am going to be using AIS CLI. But we can also use the Python SDK for the same.

1
! ais config cli set cluster.url=http://localhost:51080
2
3
# create bucket using AIS CLI
4
! ais create caltech256
5
6
# put the downloaded dataset in the created AIS bucket
7
! ais put -r -y <path_to_dataset> ais://caltech256/

OUTPUT:
“ais://caltech256” created (see https://github.com/NVIDIA/aistore/blob/main/docs/bucket.md#bucket-properties)
Files to upload:
EXTENSION COUNT SIZE
1 3.06KiB
.jpg 30607 1.08GiB
TOTAL 30608 1.08GiB
PUT 30608 objects to “ais://caltech256”

Preloaded dataset

The following assumes that AIS cluster is running and one of its buckets contains Caltech-256 dataset.

1
# list of prefixes which contain data
2
image_prefix = ['ais://caltech256/']
3
4
# Listing all files starting with these prefixes on AIStore 
5
dp_urls = AISFileLister(url="http://localhost:51080", source_datapipe=image_prefix)
6
7
# list first 5 obj urls
8
list(dp_urls)[:5]

OUTPUT:
[‘ais://caltech256/002.american-flag/002_0001.jpg’,
‘ais://caltech256/002.american-flag/002_0002.jpg’,
‘ais://caltech256/002.american-flag/002_0003.jpg’,
‘ais://caltech256/002.american-flag/002_0004.jpg’,
‘ais://caltech256/002.american-flag/002_0005.jpg’]

1
# loading data using AISFileLoader
2
dp_files = AISFileLoader(url="http://localhost:51080", source_datapipe=dp_urls)
3
4
# check the first obj
5
url, img = next(iter(dp_files))
6
7
print(f"image url: {url}")
8
9
# view the image
10
# Image(data=img.read())

OUTPUT:
image url: ais://caltech256/002.american-flag/002_0001.jpg

1
def collate_sample(data):
2
    path, image = data
3
    dir = os.path.split(os.path.dirname(path))[1]
4
    label_str, cls = dir.split(".")
5
    return {"path": path, "image": image, "label": int(label_str), "cls": cls}

1
# passing it further down the pipeline
2
for _sample in Mapper(dp_files, collate_sample):
3
    pass

Remote cloud buckets

AIStore supports multiple remote backends. With AIS, accessing cloud buckets doesn’t require any additional setup assuming, of course, that you have the corresponding credentials (to access cloud buckets).

For the following example, AIStore must be built with --gcp build tag.

--gcp, --aws, and a number of other build tags is the mechanism we use to include optional libraries in the build.

1
# list of prefixes which contain data
2
gcp_prefix = ['gcp://webdataset-testing/']
3
4
# Listing all files starting with these prefixes on AIStore 
5
gcp_urls = AISFileLister(url="http://localhost:51080", source_datapipe=gcp_prefix)
6
7
# list first 5 obj urls
8
list(gcp_urls)[:5]

OUTPUT:
[‘gcp://webdataset-testing/coco-train2014-seg-000000.tar’,
‘gcp://webdataset-testing/coco-train2014-seg-000001.tar’,
‘gcp://webdataset-testing/coco-train2014-seg-000002.tar’,
‘gcp://webdataset-testing/coco-train2014-seg-000003.tar’,
‘gcp://webdataset-testing/coco-train2014-seg-000004.tar’]

1
dp_files = AISFileLoader(url="http://localhost:51080", source_datapipe=gcp_urls)

1
for url, file in dp_files.load_from_tar():
2
    pass

References

• AIStore
• AIStore Blog
• AIS CLI
• AIStore Cloud Backend Providers
• AIStore Documentation
• AIStore Python SDK
• Caltech 256 Dataset
• Getting started with AIStore
• Microsoft COCO Dataset