TensorFlow分割模型多并行预测结果异常问题求助

我训练了一个图像分割模型（输入图像尺寸为(32,32,3)，输出掩码尺寸为(32,32)），推理阶段加载了尺寸为(62500, 32, 32, 3)的图像数组tiles，原本采用单张循环预测的方式：

masks = [] 
for k, element in enumerate(the_image_array):
        the_img = np.asarray(np.expand_dims(element, 0))[-1, -1, :, :]
        pred = model.predict(the_img[np.newaxis, :, :, :])[0]
        mask = tf.where(pred > 0.5, 255, 0)
        masks.append(mask)

为提升效率，我尝试用多进程实现并行预测，代码如下：

import tensorflow as tf
import numpy as np
import os
from tensorflow.keras.models import load_model
from itertools import chain
from tensorflow.keras import backend as K
import multiprocessing
from multiprocessing import Pool

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

multiprocessing.set_start_method('spawn', force=True)


model = load_model('./model.h5',
                   custom_objects={"K": K})
     

def resize_and_rescale(image):
    image = tf.image.resize(image, 
                            (32, 32),
                            preserve_aspect_ratio=True)
    image /= 255.0
    return image
    
def prepare(ds):
    ds = ds.map(resize_and_rescale)
    return ds

def _apply_df(data):
    img = np.asarray(np.expand_dims(data, 0))[-1,-1, :, :]
    print(img.shape)
    pred = model.predict(img[np.newaxis,  :, :, :], verbose=2)[0]
    
    #pred = model.predict(data)[0]
    mask = tf.where(pred[:, :, -1] > 0.5, 255, 0)
    return mask

def apply_by_multiprocessing(data, workers):

    pool = Pool(processes=workers)   
    #result = pool.map(_apply_df, np.array_split(list(data.as_numpy_iterator()), workers))
    result = pool.map(_apply_df, data.batch(np.ceil(len(data) / workers)))
    pool.close()
    return list(result)


def after_prepare(data):
    tens_data = tf.data.Dataset.from_tensor_slices(data)
    tens_data = prepare(tens_data)
    return tens_data


def main():
    
    tiles = np.load('tiles.npy')
    print(len(tiles))
    print(tiles[0].shape)
    prep = after_prepare(tiles)
    print(len(prep))
    
    masks = apply_by_multiprocessing(prep, workers=4)
      
    masks_flatten = list(chain.from_iterable(masks))
    print(len(masks_flatten), masks_flatten[0].shape)  # 
    
    return masks_flatten
   
    
    
if __name__=="__main__":
    masks_flatten = main()

但运行后得到的masks_flatten长度为128，元素形状为(32,)，完全不符合预期的62500个形状为(32,32)的掩码。

现在更新需求：希望实现无需内部循环的并行预测，逻辑等价于以下代码：

def _apply_df(data):
    results = []
    for el in data:
        pred = model.predict(el[np.newaxis, :, :, :], verbose=2)[0]
        mask = tf.where(pred[:, :, -1] > 0.5, 255, 0)
        results.append(mask)
        
    return results

内容的提问来源于stack exchange，提问作者George