启用GPU运行tf.agents时卷积层报错问题咨询

tf.agents 原生支持GPU环境下的卷积层运算，不存在框架层面的功能限制。Colab GPU运行时已经预装了卷积加速所需的CUDA、cuDNN依赖，不需要额外安装其他组件。

CPU环境运行正常、GPU环境报错的核心原因是两类运行时的执行逻辑差异，触发了代码中的隐式问题，常见诱因如下：

• 卷积维度顺序不匹配：CPU运行时对张量维度格式容错性较高，即使输入观测的张量维度顺序（NHWC/NCHW）和卷积层预期不匹配，也会自动做隐式转换完成计算；GPU调用cuDNN卷积加速算子时对维度格式校验严格，如果没有给卷积层显式指定data_format参数，很容易因为维度顺序不匹配抛出运行错误。
• 环境版本不兼容：Colab切换GPU运行时后，默认预装的TensorFlow与tf-agents版本经常存在版本差，卷积层相关的层封装接口在版本迭代中做过调整，CPU运行时走通用计算内核不会触发兼容bug，GPU调用原生算子时就会触发异常。这类问题可以先执行!pip install --upgrade tensorflow tf-agents，重启运行时后再测试，绝大多数版本兼容问题都可以通过这个方式解决。
• 张量设备放置错误：自定义叠加卷积层时，如果没有将卷积层和输入张量统一放置到同一GPU设备上，或是预处理阶段的卷积输入留在CPU内存、没有做显式设备拷贝，GPU执行卷积运算时无法访问对应内存块的张量数据就会报错；CPU环境下不存在设备内存隔离，因此不会触发这类问题。

补充说明：tf-agents官方所有基于视觉输入的强化学习示例（包括Atari环境下的标准DQN实现），默认都支持GPU运行卷积层结构，不存在功能层面的支持缺陷。

内容的提问来源于stack exchange，提问作者José Luis Neves