Using AIStore with built-in prototype datasets

[austinmw]

📚 The doc issue

Hi, I'd like to be able to download prototype datasets to my AIStore cluster and use AISFileLister, AISFileLoader.

Also, how can I gain access to CifarFileReader and associated functions if my resource is AIStore? For example, currently I did the following:

• created an AIStore cluster
• downloaded data (cifar10) locally
• put data in aistore
• Created initial dp with AISFileLister, AISFileLoader
• copied and pasted all the functions and classes defined in the cifar file to be able to use Filter(dp, _is_data_file), Mapper(dp, _unpickle), CifarFileReader(dp, labels_key="labels"), Mapper(dp, _prepare_sample)

---

My final script is below, but I assume there must be a way to make it more concise?

from torchdata.datapipes.iter import AISFileLister, AISFileLoader

# list of prefixes which contain data
image_prefix = ["ais://cifar10/"]

# Listing all files starting with these prefixes on AIStore
dp_urls = AISFileLister(url="http://localhost:51080", source_datapipe=image_prefix)

# list obj urls
print(list(dp_urls))

# loading data using AISFileLoader
dp = AISFileLoader(url="http://localhost:51080", source_datapipe=dp_urls)


import io
import pathlib
import pickle
from typing import Any, BinaryIO, cast, Dict, Iterator, List, Optional, Tuple, Union

import numpy as np
import torch
from torchdata.datapipes.iter import TarArchiveLoader, Filter, IterDataPipe, Mapper
from torchvision.prototype.datapoints import Image, Label
from torchvision.prototype.datasets.utils._internal import (
    hint_sharding,
    hint_shuffling,
    path_comparator,
    read_categories_file,
)

class CifarFileReader(IterDataPipe[Tuple[np.ndarray, int]]):
    def __init__(self, datapipe: IterDataPipe[Dict[str, Any]], *, labels_key: str) -> None:
        self.datapipe = datapipe
        self.labels_key = labels_key

    def __iter__(self) -> Iterator[Tuple[np.ndarray, int]]:
        for mapping in self.datapipe:
            image_arrays = mapping["data"].reshape((-1, 3, 32, 32))
            category_idcs = mapping[self.labels_key]
            yield from iter(zip(image_arrays, category_idcs))

categories = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

split = "train"

def _is_data_file(data: Tuple[str, Any]) -> bool:
    path = pathlib.Path(data[0])
    return path.name.startswith("data" if split == "train" else "test")

def _unpickle(data: Tuple[str, io.BytesIO]) -> Dict[str, Any]:
    _, file = data
    content = cast(Dict[str, Any], pickle.load(file, encoding="latin1"))
    file.close()
    return content

def _prepare_sample(data: Tuple[np.ndarray, int]) -> Dict[str, Any]:
    image_array, category_idx = data
    return dict(
        image=Image(image_array),
        label=Label(category_idx, categories=categories),
    )


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

dp = TarArchiveLoader(dp)
dp = Filter(dp, _is_data_file)
dp = Mapper(dp, _unpickle)
dp = CifarFileReader(dp, labels_key="labels")
dp = hint_shuffling(dp)
dp = hint_sharding(dp)
dp = Mapper(dp, _prepare_sample)
dp = dp.batch(8).collate()

for x in dp:
    print(x.keys())
    images = x['image'].to(device)
    labels = x['label'].to(device)
    print(images.shape)
    break

Suggest a potential alternative/fix

Demonstrating how to use the prototype dataset classes if you have your own data stored in a matching format

cc @pmeier @bjuncek

[NicolasHug]

Hi @austinmw ,

Sorry, I don't really understand what the problem is. Is there any failure?

My current understanding of your question is "How do I make datapipes work with AIStore" but this is probably a question that should be directed to the datapipe or AIStore projects directly.

[pmeier]

IIUC, the request is the same as #7147 but for CIFAR? To summarize, you want a way to change the source where we are loading the dataset from, rather than hardcoding the local loading, right?

Unfortunately, there currently is no way to do so and in the foreseeable future nothing in the same direction is planned. If you are ok with hacky solutions, I detailed in #7147 (comment) how I would monkey patch our logic to get custom local loading going. Replacing this with custom loading from another source should also be doable. Let me know if you encounter issues and I might be able to point you in the right direction.

[austinmw]

@pmeier Thanks, yeah I guess it's essentially the same request. I'd like to use everything from the CIFAR pipeline except for where the dataset is loaded from, without needing to copy/paste all of the code. For example just being able to import the CifarFileReader would be nice.

I was able to get this working using the method you described, but am hoping a cleaner solution becomes available in the future. It does seem less clean than ideal to need to copy/paste all that code to load local cifar, coco, etc. data.