在使用Tensorflow 分布式训练出现的Non-OK-status: GpuLaunchKernel问题

项目场景：

使用Tensorflow的keras进行模型构建，然后使用tf.distribute.MirroredStrategy进行多gpu训练

---

问题描述

在调试完其他部分然后开始训练，执行model.fit语句时报出了下面的错误，之后程序就停止运行了。

Non-OK-status: GpuLaunchKernel( SwapDimension1And2InTensor3UsingTiles<T, NumThreads, TileLongSide, TileShortSide>, total_tiles_count, NumThreads, 0, d.stream(), input, input_dims, output) status: Internal: invalid configuration argument
Aborted (core dumped)
	}

但在使用单gpu进行训练时未发现有错误。

---

原因分析：

batch_size和gpu数量关系不匹配导致的，在使用多个gpu训练时，每个gpu上的负载的数据量为batch_size / gpu_num，因此batch_size要是使用gpu数量的整数倍。我这里之前是用单gpu进行训练的，batch_size=1，然后用的是2个gpu进行训练，batch_size没有更改导致的。

---

解决方案：

将batch_size改为使用的gpu数量的整数倍即可。我的程序中就将batch_size修改为batch_size=2完美解决。

参考内容：Tensorflow mirrored strategy error: Non-OK-status: GpuLaunchKernel( SwapDimension1And2InTensor3UsingTiles)