1. Simplify Data
Conversion from
Spark to Deep
Learning
Liang Zhang
Software Engineer @ databricks

2. About Me
▪ Machine Learning Team
@ Databricks
▪ Master in Carnegie
Mellon University Liang Zhang
linkedin.com/in/liangz1/

3. Agenda
▪ Why should we care
about data conversion
between spark and deep
learning frameworks?
▪ Pain points
▪ Overview of the Spark
Dataset Converter
▪ Demo
▪ Best Practices

4. Spark
DataFrame
Motivation: Data Conversion from Spark to DL
TensorFlow
PyTorch
?
• Images from driving camera: Detect traffic lights
• Large amount of data - TBs
• New images arriving every day
• Data cleaning and labeling
• Train the model with all available data and periodically re-train with new data
• Predict the label of new images

5. Pain points: Data Conversion from Spark to Deep
Learning frameworks

6. Pain points: Data Conversion from Spark to DL
• Single-node training:
• Collect a sample of data to the driver in a pandas DataFrame
• Distributed training:
• Save the Spark DataFrame to TFRecords files and load TFRecords using
TensorFlow
• Save the Spark DataFrame to parquet files and write your custom PyTorch
DataLoader to load the partitions

7. Pain points: Data Conversion from Spark to DL
• Single-node training:
• Collect a sample of data to the driver in a pandas DataFrame
• Distributed training:
• Save the Spark DataFrame to TFRecords files and parse the serialized data
in TFRecords using TensorFlow
• Save the Spark DataFrame to parquet files and write your custom PyTorch
DataLoader to load the partitions
• Hard to migrate from single-node to distributed training
• Many lines of extra code to save, load and parse intermediate
files

8. Overview of the Spark Dataset Converter

9. Spark
DataFrame
Spark Dataset Converter API Overview
TensorFlow
Dataset
PyTorch
DataLoader
Spark
Dataset
Converter
from petastorm.spark import make_spark_converter
converter = make_spark_converter(df)
with converter.make_tf_dataset() as dataset:
tf_model.fit(dataset)
with converter.make_torch_dataloader() as dataloader:
train(torch_model, dataloader)

10. Spark Dataset Converter API
HDFS/DBFS
Spark
DataFrame
tf.data.Dataset /
torch.dataloader
Found
cached
parquet file?
Cache
DataFrame in
parquet file
data.parquet
No
Yes Load cached
parquet file with
petastorm
ETL Training

11. Spark Dataset Converter Features
▪ Recognize cached Spark
DataFrame by checking
the analyzed query plan
▪ Automatic cache cleaning
at program exit
• Change two arguments
to migrate your data
loading code from
single-node to
distributed setting
• Easy migration to distributed
• Cache intermediate files
• Convert MLlib vectors to
1D arrays automatically
• MLlib vector Handling

12. How to use the Spark Dataset Converter API?
(demo)

13. Demo notebooks
• Image Classification
• Spark to TensorFlow Dataset
• https://docs.databricks.com/_static/notebooks/deep-learning/petastorm-spark-converter-tenso
rflow.html
• Spark to PyTorch DataLoader
• https://docs.databricks.com/_static/notebooks/deep-learning/petastorm-spark-converter-pytor
ch.html

14. Best Practices

15. Best Practices with Spark Dataset Converter
• Image data decoding and preprocessing
• Decode image bytes and preprocess in TransformSpec, not in Spark
• Spark -> TransformSpec -> Dataset.map -> in the model (GPU)
• Generate infinite batches using num_epochs=None
• In distributed training, to guarantee that every worker get exactly the same
amount of data.
• Manage the lifecycle of cache data
• On local laptop or in a scheduled job on Databricks, the cache files will be
automatically deleted when the python process exits.
• In Databricks notebook, we recommend configuring lifecycle rules for the
underlying S3 buckets storing the cache files.

16. Feedback
Your feedback is important to us.
Don’t forget to rate and review the sessions.