<jupyter_start><jupyter_text>Creating a Petastorm MNIST datasetIn this notebook we are going to create a Petastorm dataset from the famous MNIST dataset. Compared to ImageNette it has the advantage of being easily available through PyTorch. It is also considerably smaller which makes it easier to experiment with.<jupyter_code>from hops import hdfs
import numpy as np
from torchvision.datasets import MNIST
path = hdfs.project_path() + "Resources/Petastorm"<jupyter_output><empty_output><jupyter_text>Downloading the dataset with torchvisionTorchvision provides a simple interface to download the MNIST dataset. Note that the download prior to version 0.9.1 is broken! If you have issues with this, please upgrade your installation to the latest version. For other workarounds, see [here](https://stackoverflow.com/questions/66577151/http-error-when-trying-to-download-mnist-data).<jupyter_code>path = hdfs.project_path() + "DataSets/MNIST"
train_dataset = MNIST(path, download=True)
test_dataset = MNIST(path, download=True, train=False)<jupyter_output>Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 503: Service Unavailable

Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/MNIST/MNIST/raw/train-images-idx3-ubyte.gz
Extracting hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/MNIST/MNIST/raw/train-images-idx3-ubyte.gz to hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/MNIST/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Failed to download (trying next):
HTTP Error 503: Service Unavailable

Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to hdfs://rpc.[...]<jupyter_text>Setting up the petastorm dataset generationNow that we have our dataset, creating the petastorm dataset is exactly the same as with ImageNette. Note that for distributed training you need an even dataset. If your dataset is not even (meaning that each node sees the same amount of examples) you can increase the number of parquet files in order to allow for a more fine grained distribution among nodes.<jupyter_code>from petastorm.codecs import CompressedImageCodec, NdarrayCodec, ScalarCodec
from petastorm.etl.dataset_metadata import materialize_dataset
from petastorm.unischema import Unischema, UnischemaField, dict_to_spark_row
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType


MNISTSchema = Unischema('ScalarSchema', [
   UnischemaField('image', np.uint8, (1,28,28), NdarrayCodec(), False),
   UnischemaField('label', np.int8, (), ScalarCodec(IntegerType()), False)])

def row_generator(idx, dataset):
    img, label = dataset[idx]
    return {'image': np.expand_dims(np.array(img, dtype=np.uint8), axis=0), 'label': label}


def generate_MNIST_dataset(output_url, dataset):
    rowgroup_size_mb = 1
    rows_count = len(dataset)
    parquet_files_count = 100
    
    sc = spark.sparkContext
    # Wrap dataset materialization portion. Will take care of setting up spark environment variables as
    # well as save petastorm specific metadata
    with materialize_dataset(spark, output_url, MNISTSchema, rowgroup_size_mb):
        rows_rdd = sc.parallelize(range(rows_count))\
            .map(lambda x: row_generator(x, dataset))\
            .map(lambda x: dict_to_spark_row(MNISTSchema, x))

        spark.createDataFrame(rows_rdd, MNISTSchema.as_spark_schema()) \
            .repartition(parquet_files_count) \
            .write \
            .mode('overwrite') \
            .parquet(output_url)<jupyter_output><empty_output><jupyter_text>Generating the datasetNow that everything is set up, we can define our output paths and generate the datasets.<jupyter_code>train_path = hdfs.project_path() + "DataSets/MNIST/PetastormMNIST/train_set"
test_path = hdfs.project_path() + "DataSets/MNIST/PetastormMNIST/test_set"
generate_MNIST_dataset(train_path, train_dataset)
generate_MNIST_dataset(test_path, test_dataset)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09376432828977534, 'fraction_numerical': 0.012379642365887207, 'mean_word_length': 5.26867816091954, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '213568', 'license': 'Apache-2.0', 'n_tokens_mistral': 1452, 'n_tokens_neox': 1383, 'n_words': 341, 'path': 'examples/MNIST-to-petastorm.ipynb', 'repo_name': 'ssheikholeslami/maggy', 'size': '4362'}