Python 在tensorflow中，是否有可能在神经网络内将数据分成批次？

我有大量的图像要通过分割卷积神经网络。为了能够在不耗尽GPU资源的情况下处理整个图像，我需要将输入图像分割为块或补丁，分批传递它们，然后通过将输出补丁正确重塑为其原始位置来重建全尺寸图像

在网络结束时，我想执行一个考虑到完整图像的操作。但是为了这个，我需要对整个图像进行分割

因此，这些步骤是：

将输入图像重塑为不重叠的“平铺”

将瓷砖分批传递，并在最后聚集

在各个批次中对整个图像进行处理后，重建原始图像大小分割

对完整图像分割执行操作（识别感兴趣区域和边界框）

有没有办法将所有4个步骤都包含在神经网络中？我可以创建一个自定义层，将其重塑为瓷砖，然后再重塑回图像，但我想不出一种方法也可以批量传递输入，并在最后一层上重建原始图像之前在最后进行聚集

这可能吗？如果是，如何做到这一点

我知道tensorflow.data.Dataset API具有批处理操作，但不确定如何在继续下一层之前在另一端“收集”批处理

另一个至少在概念上有意义的选择是拥有一个“嵌套”网络。顶部网络只是将一个重塑输入到tiles层中，然后将数据分割成批，这些批被送入执行分割的“内部”网络。内部网络完成所有批处理后，外部网络将获取结果并将其重塑为完整图像大小。可行吗