CUDA out of memory. Tried to allocate 896.00 MiB. GPU 0 has a total capacity of 22.18 GiB of which 4

torch.OutOfMemoryError 错误表示 GPU 内存不足，无法为当前操作分配所需的内存。这通常是因为在加载模型或处理数据时，所需的 GPU 内存超过了可用的内存。以下是一些可能的解决方案：

1. 减小批处理大小

• 调整批处理大小：尝试减小输入的批处理大小（batch size），这样可以减少每次操作所需的 GPU 内存。

2. 释放显存

• 重启进程：如果在之前的运行中有未释放的显存，重启 Python 进程可能有助于释放内存。
• 清理显存：在使用 PyTorch 时，可以使用 torch.cuda.empty_cache() 手动清理未使用的显存。

3. 使用内存管理设置

• 设置内存配置：可以通过设置环境变量来优化内存管理：

  export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"
  

  或者使用你提到的配置：

  export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
  

4. 检查模型和数据

• 模型大小：确认所使用的模型是否过大，考虑选择一个较小的模型或版本。
• 数据预处理：确保输入数据经过适当处理，避免因过大的输入导致显存不足。

5. 分布式训练

• 分布式训练：如果可能，考虑在多个 GPU 之间分配负载，通过分布式训练（如使用 PyTorch 的 DistributedDataParallel）来管理显存使用。

6. 模型精度

• 使用混合精度训练：考虑使用自动混合精度 (AMP) 以减少显存占用：

  from torch.cuda.amp import autocast
  
  with autocast():
      # 运行模型的代码
  

7. 监控 GPU 使用

• 监控工具：使用 nvidia-smi 命令实时监控 GPU 的使用情况，查看哪些进程在占用显存。

如果在尝试上述方法后问题依然存在，可能需要更详细地审查模型的具体实现和使用方式。请根据实际情况调整参数和配置，确保在每次修改后进行测试。