PyTorch Geometric中ModuleDict与TorchScript的兼容性问题分析

问题背景

在使用PyTorch Geometric(简称PyG)进行图神经网络开发时，开发者发现当尝试将包含tgnn.module_dict.ModuleDict的自定义网络转换为TorchScript格式时会出现兼容性问题。这个问题在PyG 2.5.3版本中表现为类型错误，而在较早的2.3版本中却能正常工作。

问题现象

当开发者使用PyTorch Geometric的ModuleDict实现一个简单的卷积网络并尝试通过torch.jit.script进行脚本化时，会收到以下错误信息：

TypeError: 'set' object in attribute 'ModuleDict.CLASS_ATTRS' is not a valid constant.

相比之下，如果使用PyTorch原生的nn.ModuleDict则不会出现这个问题。

技术分析

TorchScript的限制

TorchScript对Python代码的转换有其特定的限制和要求。它要求模块属性必须是以下类型之一：

1. 一个nn.ModuleList
2. 基本类型值(bool, float, int, str, NoneType等)
3. 设备或数据类型(torch.device, torch.layout, torch.dtype)
4. 上述类型的列表或元组

PyG ModuleDict的实现差异

PyTorch Geometric的ModuleDict实现与原生PyTorch的实现存在差异，特别是在处理类属性CLASS_ATTRS时。PyG版本使用了Python的set类型来存储这些属性，而TorchScript不支持将set类型作为模块的常量属性。

版本变化的影响

在PyG 2.3版本中能够正常工作，而在2.5.3版本中出现问题，这表明在版本迭代过程中，ModuleDict的实现发生了变化，可能涉及内部属性的存储方式调整。

解决方案

PyG开发团队已经注意到这个问题，并在后续版本中进行了修复。修复方案主要涉及修改ModuleDict的实现，确保其属性类型符合TorchScript的要求。

开发者建议

对于遇到类似问题的开发者，可以采取以下临时解决方案：

1. 优先使用PyTorch原生的nn.ModuleDict，除非确实需要PyG特有的功能
2. 检查PyG的更新日志，确认该问题是否已在最新版本中修复
3. 如果必须使用PyG的ModuleDict，可以考虑回退到2.3版本

总结

PyTorch Geometric作为PyTorch的扩展库，在提供图神经网络专用功能的同时，也需要保持与PyTorch核心功能的兼容性。这个案例展示了在模块属性设计时考虑TorchScript兼容性的重要性，也为开发者提供了处理类似兼容性问题的思路。