AIStore & ETL: Using AIS/PyTorch connector to transform ImageNet (post #2)

The goal now is to deploy our first ETL and have AIStore run it on each storage node, harnessing the distributed power (and close to data - meaning, fast). For the problem statement, background and terms, please see the previous post:

• AIStore & ETL: Introduction

To quickly get to the point, we’ll assume that an instance of AIStore has been already deployed on Kubernetes.

Check out our dedicated ais-k8s repository for the multiple easy ways to accomplish Kubernetes deployments. For a quick local deployment without the operator, refer to deploy/dev/k8s.

We’ll be using PyTorch’s torchvision to transform ImageNet dataset - as illustrated:

Also, in the examples below you’ll notice ais command. That’s AIStore’s CLI tool providing unmatched (well, almost) convenience and ease-of-use. Most of the time, though, we’ll show the equivalent curl.

For a variety of practical reasons, curl proves to be handly in use cases - in bash and Python scripts on the client side. Big part of that is its (i.e., the curl’s) ubiquitous nature. Great tool, overall.

The Dataset

The dataset we have here is derived from the original ImageNet and is only slightly different. Its training part exists under a train/ directory, validation - under val/. Each *.jpg image has a corresponding *.cls object with the corresponding class number. Image are assigned to one of the 1000 classes; classes are represented as integers in the [0 - 999] range:

1
$ ais ls ais://imagenet
2
NAME             SIZE
3
train/0000353.cls     1B
4
train/0000353.jpg     17.88KiB
5
...
6
val/1280048.cls       3B
7
val/1280048.jpg       82.78KiB

A random (non-transformed) image from the dataset (and again, notice ais CLI usage):

1
$ ais get ais://imagenet/train/0278350.jpg
2
$ open 0278350.jpg

And the associated class:

1
$ ais object cat ais://imagenet/train/0278350.cls
2
217

The Plan and the Code

The plan, essentially, is two-fold:

1. Deploy provided transformation code (called code.py below) as ETL K8s container aka transformer.
2. Drive transformer from the PyTorch-based client to transform requested objects (shards) as required.

In the end, each image from the dataset, before it reaches the model, goes through a series of the following (code.py) transformations:

1
# `code.py`:
2
import io, sys
3
import torch
4
from PIL import Image
5
from torchvision import transforms
6
7
def img_to_bytes(img):
8
    buf = io.BytesIO()
9
    img = img.convert('RGB')
10
    img.save(buf, format='JPEG')
11
    return buf.getvalue()
12
13
preprocessing = transforms.Compose([
14
    transforms.RandomResizedCrop(224),
15
    transforms.RandomHorizontalFlip(),
16
    transforms.ToTensor(),
17
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
18
    transforms.ToPILImage(),
19
    transforms.Lambda(img_to_bytes),
20
])
21
22
if __name__ == "__main__":
23
    input_bytes = sys.stdin.buffer.read()
24
    image = Image.open(io.BytesIO(input_bytes)).convert('RGB')
25
    processed_bytes = preprocessing(image)
26
    sys.stdout.buffer.write(processed_bytes)

Initializing

We will use python3 (python:3.9) runtime to install torch and torchvision packages.

To make sure that code.py (above) can have its imports, the following (deps.txt) dependencies must be installed:

torch==1.6.0
torchvision==0.7.0

With transforming code and dependencies covered, we are now fully ready to initialize ETL in the cluster:

1
$ ais etl init code \
2
  --name my-first-etl \
3
  --from-file code.py \
4
  --deps-file deps.txt \
5
  --runtime python3 \
6
  --comm-type io://

Notice the “elements” of this ais command:

• user-given name of the specific (code.py) transformation;
• the dependencies;
• the aforementioned runtime, and, finally -
• --comm-type (“communication type”) option briefly already mentioned - and we’ll discuss it in-depth in our future postings.

Two ways to check that ETL is up and running: ais CLI (“way”) and kubectl:

$ ais etl list
NAME
my-first-etl
$ kubectl -n ais get pods | grep ‘my-first-etl’
ais   my-first-etl-iacjhrvc                            1/1     Running   0          50s

Recap what we just did:

1. We prepared Python3 code (code.py) and provided dependencies (deps.txt) to run it.
2. We started transformer in the AIS cluster to, subsequently, augment images from the ImageNet dataset.

Transforming a single object

To get a single object, we will use the AIStore’s CLI, and we will show equivalent curl commands which can be useful in certain situations. Those commands are convenient for quick testing and specific solutions which are written e.g. in bash.

The original image:

1
$ ais get ais://imagenet/train/0278350.jpg 0278350.jpg
2
3
# Equivalent `curl`:
4
$ curl -O 0278350.jpg https://aistore/v1/objects/imagenet/train/0278350.jpg

and, the image after the transformation:

1
$ ais etl object my-first-etl ais://imagenet/train/0278350.jpg 0278350.jpg
2
3
# Equivalent `curl`:
4
$ curl -O 0278350.jpg https://aistore/v1/objects/imagenet/train/0278350.jpg?uuid=my-first-etl

Notice that the only difference between these two curl (and, respectively, ais) commands above - is the uuid parameter referencing existing and available (my-first-etl) transformation that we have just previously initialized.

Post-transform 0278350.jpg image:

AIS/PyTorch connector

So far we have set up ETL and tried our first cluster-resident transformations. We can now start running a real training model. For this purpose, we have prepared a slightly modified version of the PyTorch ImageNet example that can be found here. The script contains training and validation code for the ImageNet dataset.

Next step is to modify the script to utilize my-first-etl transformer.

The typical code for loading ImageNet from a local directory looks like this:

1
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
2
3
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
4
       os.path.join(args.data, 'train'),
5
       transforms.Compose([
6
           transforms.RandomResizedCrop(224),
7
           transforms.RandomHorizontalFlip(),
8
           transforms.ToTensor(),
9
           normalize,
10
       ]),
11
    ),
12
    batch_size=args.batch_size, shuffle=True,
13
    num_workers=args.workers, pin_memory=True)
14
15
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
16
       os.path.join(args.data, 'val'),
17
       transforms.Compose([
18
           transforms.Resize(256),
19
           transforms.CenterCrop(224),
20
           transforms.ToTensor(),
21
           normalize,
22
       ]),
23
    ),
24
    batch_size=args.batch_size, shuffle=False,
25
    num_workers=args.workers, pin_memory=True)

Full code for the example above is also available - see ImageNet PyTorch training with dataset.ImageFolder.

In the world of PyTorch, all datasets are subclasses of torch.utils.data.Dataset with torchvision.datasets.ImageFolder being the standard for handling datasets with labeled images.

To integrate with PyTorch and offload transformations to AIStore, we introduce aistore.pytorch.Dataset - the implementation of torch.utils.data.Dataset.

With aistore.pytorch.Dataset, the example above works out as follows:

1
import aistore
2
from aistore.client import Bck
3
4
...
5
6
train_loader = torch.utils.data.DataLoader(
7
    aistore.pytorch.Dataset(
8
        "http://aistore-sample-proxy:51080", # AIS IP address or hostname
9
	Bck("imagenet"),
10
        prefix="train/", transform_id="my-first-etl",
11
        transform_filter=lambda object_name: object_name.endswith('.jpg'),
12
    ),
13
    batch_size=args.batch_size, shuffle=True,
14
    num_workers=args.workers, pin_memory=True)
15
16
val_loader = torch.utils.data.DataLoader(
17
    aistore.pytorch.Dataset(
18
        "http://aistore-sample-proxy:51080", # AIS IP address or hostname
19
	Bck("imagenet"),
20
        prefix="val/", transform_id="my-second-etl", # We skipped setting up this ETL.
21
        transform_filter=lambda object_name: object_name.endswith('.jpg'),
22
    ),
23
    batch_size=args.batch_size, shuffle=False,
24
    num_workers=args.workers, pin_memory=True)

Complete code is available here:

• ImageNet PyTorch training with aistore.pytorch.Dataset

References

1. AIStore & ETL: Introduction
2. GitHub:
   • AIStore
   • AIS/Kubernetes Operator, AIS on bare-metal, Deployment Playbooks, Helm
   • AIS-ETL containers and specs
3. Documentation, blogs, videos:
   • https://aiatscale.org
   • https://github.com/NVIDIA/aistore/tree/main/docs

PS. Note that we have omitted setting-up ETL for the validation loader - leaving it as an exercise for the reader. To be continued…