Gpu 在分割模型的多 GPU 训练中丢失

我正在尝试训练工艺 OCR 模型。我正在使用图像高达 850K 的合成文本数据集。对于每个图像，我必须创建分割groundtruth，即字符级分割和亲和力分数分割。所以这对每张图像来说都是非常繁重的预处理步骤。

因此，我不是为每个批次动态创建地面实况，而是创建包含图像的 tfrecords 并尊重地面实况分割。

因此输入图像大小为 768,768,3，该图像的真实分割为 384,384,2。

我创建了 tfrecords 分片，每个分片包含 4096 个图像及其受人尊敬的基本事实。因此，首先我尝试用 2 个分片训练模型，每个分片大小为 33GB，每个分片包含 4096 个图像及其预处理的分割基本事实。

在 keras 多 GPU 培训的帮助下，我只是在阅读 tfrecords。并从它们中批量获取图像和标签并进行训练。

但是仅仅几步之后，它就显示我的 gpu 丢失了。

但是，如果我在训练时即时生成地面实况，所有 GPU 都可以正常工作，但由于预处理繁重（生成地面实况分割），训练需要很长时间。

为了避免这种情况，我采用了这种 tfrecords 方法，其中已经存储了图像及其预处理标签。但是在训练中我面临这个 gpu 丢失问题..