RuntimeError: as_numpy_iterator() is not supported while tracing functions on tensorflow 2.x dataset on sagemaker

Viewed 322

I am preparing a sagemaker PIPE mode dataset to train a time series model on SageMaker with PIPE mode. The PipeModeDataset is a TensorFlow Dataset for reading SageMaker Pipe Mode channels. I am using an augmented manifest file which contains image location on S3 and the label each line. My model accept batches of images (512 x 512 x 1) with single label per batch as input. I thought of using the window function to bundle the images read from pipe. Please refer to following partial code for dataset generation.

def _input_fn(channel):
    """Returns a Dataset for reading from a SageMaker PipeMode channel."""
    features = {
        'image-ref': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([3], tf.int64),
    }
    
    def parse(record):
        parsed = tf.io.parse_single_example(record, features)
        image = tf.io.decode_png(parsed['image-ref'], channels=1, dtype=tf.uint8)
        image = tf.reshape(image, [512, 512, 1])
        label = parsed['label']
        return (image, label)

    ds = PipeModeDataset(channel, record_format='TFRecord', benchmark=True, benchmark_records_interval=100)
    ds = ds.map(parse)
    
    print ("PipeModeDataset print0 = " + str(ds))
    ds = ds.window(16, shift=1, drop_remainder=True)
    print ("PipeModeDataset print1 = " + str(ds))
    
    def window_func(window, label):
        window = window.batch(16, drop_remainder=True)
        label = label.batch(16, drop_remainder=True)
        
        print ("window batch is = " + str(window))
        print ("label batch is = " + str(label))
        
        window_np = np.stack(list(window.as_numpy_iterator()))
        label_np = np.stack(list(label.as_numpy_iterator())) # TODO: only get the last label
        
        return tf.data.Dataset.from_tensor_slices((window_np, label_np))
    
    ds = ds.flat_map(lambda window, label: window_func(window, label))
    ....
    ....

Getting following error at the moment. How to fix this? Recommend better ways if there is any.

PipeModeDataset print0 = <MapDataset shapes: ((512, 512, 1), (3,)), types: (tf.uint8, tf.int64)>
PipeModeDataset print1 = <WindowDataset shapes: (DatasetSpec(TensorSpec(shape=(512, 512, 1), dtype=tf.uint8, name=None), TensorShape([])), DatasetSpec(TensorSpec(shape=(3,), dtype=tf.int64, name=None), TensorShape([]))), types: (DatasetSpec(TensorSpec(shape=(512, 512, 1), dtype=tf.uint8, name=None), TensorShape([])), DatasetSpec(TensorSpec(shape=(3,), dtype=tf.int64, name=None), TensorShape([])))>
window batch is = <BatchDataset shapes: (16, 512, 512, 1), types: tf.uint8>
label batch is = <BatchDataset shapes: (16, 3), types: tf.int64>

RuntimeError: in user code:

    /opt/ml/code/train_on_pipemode.py:104 None  *
        ds = ds.flat_map(lambda window, label: window_func(window, label))
    /opt/ml/code/train_on_pipemode.py:96 window_func  *
        window_np = np.stack(list(window.as_numpy_iterator()))
    /usr/local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py:518 as_numpy_iterator  **
        raise RuntimeError("as_numpy_iterator() is not supported while tracing "

    RuntimeError: as_numpy_iterator() is not supported while tracing functions

This answer says to enable eager execution, but it is enabled in my case when I printed tf.executing_eagerly(). I am training on tensorflow 2.x.

Tensorflow version: 2.3.1
Eager execution: True
0 Answers
Related