SageMaker Script Mode + Pipe Mode

Question

I'm training in SageMaker using TensorFlow + Script Mode and currently using 'File' input mode for my data.

Has anyone figured out how to stream data using 'Pipe' data format in conjunction with Script Mode training?

Answer 1

You can import sagemaker_tensorflow from the training script as follows:

from sagemaker_tensorflow import PipeModeDataset
from tensorflow.contrib.data import map_and_batch

channel = 'my-pipe-channel-name'

ds = PipeModeDataset(channel)
ds = ds.repeat(EPOCHS)
ds = ds.prefetch(PREFETCH_SIZE)
ds = ds.apply(map_and_batch(parse, batch_size=BATCH_SIZE,
                            num_parallel_batches=NUM_PARALLEL_BATCHES))

You can find the full example here: https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/tensorflow_pipemode_example/pipemode.py

You can find documentation about sagemaker_tensorflow here https://github.com/aws/sagemaker-tensorflow-extensions#using-the-pipemodedataset