IntrinsicAvatar项目中的CUDA内存溢出问题分析与解决方案

IntrinsicAvatar项目中的CUDA内存溢出问题分析与解决方案

IntrinsicAvatar Official Implementation of IntrinsicAvatar 项目地址: https://gitcode.com/gh_mirrors/in/IntrinsicAvatar

问题背景

在使用IntrinsicAvatar项目进行人体建模训练时，部分用户遇到了CUDA内存溢出问题。该问题在使用PeopleSnapshot数据集的'male-3-casual'样本进行训练时尤为明显，表现为随着训练轮次的增加，GPU内存消耗逐渐上升，最终导致内存不足错误。

错误现象

典型的错误信息显示PyTorch尝试分配1.28GB内存失败，尽管GPU总容量为39.39GB，但已有28.72GB被分配，仅剩1.14GB可用。这种内存消耗增长模式表明可能存在内存泄漏或配置不当的问题。

原因分析

经过项目维护者的测试和验证，发现以下几个可能的原因：

1. 物理属性计算的引入：当模型开始计算物理属性时，会显著增加内存消耗，这是设计上的预期行为。

2. PyTorch-Lightning版本问题：不同版本的PyTorch-Lightning可能存在内存管理差异，特别是在验证阶段(with torch.no_grad())可能会出现内存释放不完全的情况。

3. 硬件配置不足：虽然项目在24GB显存的TITAN RTX上可以正常运行，但某些操作可能需要更多显存。

解决方案

针对上述问题，可以采取以下解决方案：

1. 调整采样参数：

   • 减少每个像素的采样数(SPP)
   • 修改configs/config.yaml文件中的相关配置
   • 对应调整models/intrinsic_avatar.py中1392-1407行的相关代码

2. 优化验证过程：

   • 通过设置trainer.val_check_interval=null来禁用训练期间的验证例程
   • 这可以避免验证阶段可能的内存泄漏问题

3. 调整分块处理大小：

   • 修改model.secondary_shader_chunk参数(例如设置为80000)
   • 这种方法可以在保持完整SPP的同时降低显存使用，但会牺牲一定的训练/推理速度

4. 环境一致性检查：

   • 确保使用Ubuntu 20.04/CentOS 7.9.2009系统
   • 使用Python 3.10、PyTorch 1.13和CUDA 11.6版本
   • 确认pytorch-lightning版本为1.9.5

最佳实践建议

对于显存有限的用户，建议采用以下配置策略：

1. 对于24GB显存及以上的GPU，可以直接使用默认配置
2. 对于显存较小的GPU，优先考虑调整secondary_shader_chunk参数
3. 在遇到内存问题时，首先尝试禁用验证过程以确认是否为内存泄漏导致
4. 保持软件环境与推荐配置一致，避免版本兼容性问题

通过合理配置这些参数，大多数用户应该能够在不同硬件配置上成功运行IntrinsicAvatar项目的人体建模训练流程。

IntrinsicAvatar Official Implementation of IntrinsicAvatar 项目地址: https://gitcode.com/gh_mirrors/in/IntrinsicAvatar