使用tf.while_循环时，使用tf.function对tensorflow数据集进行极慢的数据预处理

Question

我正在尝试预处理一个数据集，该数据集包含约20000幅图像、它们的特征向量和它们的遮罩（图像分割）。由于数据的大小，我无法执行内存中的预处理

我使用此函数加载数据：

@tf.function
def get_mask_cum_featvecs_dataset(files_paths):
    
    paths_ds = tf.data.Dataset.list_files(files_paths)
    ds = paths_ds.interleave(lambda tfrec_path: tf.data.TFRecordDataset(tfrec_path).map(
                                read_tfrec, num_parallel_calls=AUTO),
                             num_parallel_calls=AUTO)
    
    return ds

预处理包括获取数据集中存在的所有图像和输入图像（由用户提供）之间的欧几里德距离，然后将所有欧几里德距离除以存在的所有欧几里德距离中的最大欧几里德距离。 为此，我使用以下给定函数：

@tf.function
def euclidean_dist(comp_img_features, img_featvec):
    euclidean_dist = tf.sqrt(tf.reduce_sum(tf.square(img_featvec-comp_img_features), 0))
    return euclidean_dist

@tf.function
def preprocess(img_featvec, img_mask, mask_cum_featvecs_paths):
   
    tensor_img_featvec = img_featvec
    tensor_img_mask = img_mask
    tensor_mask_cum_featvecs_paths = mask_cum_featvecs_paths
    
    # the below given statements are used when this function
    # is not wrapped by tf.function decorator:
    #tensor_img_featvec = tf.constant(img_featvec)
    #tensor_img_mask = tf.constant(img_mask)
    #tensor_mask_cum_featvecs_paths = tf.constant(mask_cum_featvecs_paths)
    
    compare_img_ds = get_mask_cum_featvecs_dataset(tensor_mask_cum_featvecs_paths)
    
    get_img_euclidean_dists = partial(euclidean_dist, img_featvec=tensor_img_featvec)
    
    print("Calculating Euclidean Distances......")
    features_n_eucl_ds = compare_img_ds.map(lambda image_name, image, featvec, mask: (image_name, image, featvec, mask,
                                                                           get_img_euclidean_dists(featvec)),
                                 num_parallel_calls=AUTO)


    print("Finding Maximum Euclidean Distance........")
    num_of_loops = 2
    eucl_val_ds = features_n_eucl_ds.map(lambda image_name, image, featvec, mask, eucl_val: eucl_val,
                                         num_parallel_calls=AUTO)
   
    def dataset_reduce_max(ds, i):
        ds = ds.batch(5000)
        ds = ds.map(lambda eucl_vals: tf.math.reduce_max(eucl_vals),
                    num_parallel_calls=AUTO)
        i += 1
        return ds, i
    
    def loop_cond(ds, i):
        return tf.less(i, num_of_loops)
    
    eucl_val_ds, _ = tf.while_loop(loop_cond,
                                   dataset_reduce_max,
                                   [eucl_val_ds, tf.constant(0)])
    
    max_eucl_dist = tf.constant(0, dtype=tf.float32)

    #---------(I)---------#
    for item in eucl_val_ds.take(1):
        max_eucl_dist = item

    ratio_eucl_ds = features_n_eucl_ds.map(
                        lambda image_name, image, featvec, mask, eucl_val: (image_name,
                                                                            image,
                                                                            featvec,                                                                                             
                                                                            mask,                                                                                            
                                                                            eucl_val/max_eucl_dist),
                                           num_parallel_calls=AUTO)


    return ratio_eucl_ds #----------(II)-----------#

在tensorflow数据集中，没有用于查找特定数据流中最大元素的内置函数（此处，数据流=图像名称、图像、特征、掩码、eucl_距离） 为了找到欧几里德距离的最大值，我创建了一个单独的包含欧几里德距离的tf.dataset，并使用了数据集缩减策略。我创建了大小为5000的批，并返回其中存在的最大元素。使用此策略，我可以在2次迭代中有效地获得最大欧氏距离

问题：

当我使用tf.function decorator时，整个函数将执行得非常快，然后卡在“返回比率”（I）上。如果我不使用tf.function decorator，那么这个函数就会被困在for循环（II）中。执行代码大约需要12分钟。当我随机将任何值放入“max_eucl_dist”中时，预处理函数将无缝执行。因此，我想问题在于我试图为“max\u eucl\u dist”提取值的方式