tfds.as_numpy() 在 map function 后崩溃

Question

我尝试使用 map function 将图像旋转 20 度，并读取图像数据以确认其是否正常工作如下。

环境：

• 谷歌合作实验室
• Tensorflow 2.2

代码：

import tensorflow_datasets as tfds
import tensorflow as tf
import tensorflow_addons as tfa
import numpy as np
MODE_AUTOTUNE = tf.data.experimental.AUTOTUNE
batch_size = 128

def map_func(image, label):
    image = tf.cast(image, tf.float32) / 255.
    angle = 20. * 3.1415926 / 180.0
    image = tfa.image.rotate(images=image, angles=angle, interpolation='NEAREST')
    return image, label

## ... load x_train (images), y_train (labels) ... ##
train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_data = train_data.shuffle(len(x_train))    
train_data = train_data.map(map_func=map_func, num_parallel_calls=MODE_AUTOTUNE)
train_data = train_data.batch(batch_size)
train_data = train_data.prefetch(buffer_size=MODE_AUTOTUNE) 

datas = tfds.as_numpy(train_data)
for data in datas:
    print(data[0])
    ...
    ## try to plot datas ##

它因未知原因而坠毁。 但是，当我删除以下代码时，它运行良好。

    image = tfa.image.rotate(images=image, angles=angle, interpolation='NEAREST')

我不知道。 map function 是否运行良好，但 tfds.as_numpy() 无法读取？