DeepKE项目NER模块Few-Shot训练中的CUDA环境问题分析与解决

问题背景

在使用DeepKE项目的命名实体识别(NER)模块进行Few-Shot训练时，部分Windows用户遇到了"Torch not compiled with CUDA enabled"的错误。这个问题主要出现在使用CPU环境进行模型训练的场景下，虽然配置文件中已经明确指定了使用CPU设备，但程序仍然尝试调用CUDA功能。

错误现象分析

当用户在Windows系统下运行DeepKE的NER Few-Shot训练脚本时，会出现以下关键错误信息：

Traceback (most recent call last):
  File "run.py", line 122, in main
    trainer.train()
  ...
  File "C:\...\model.py", line 270, in get_prompt
    input_tokens = self.prompt_inputs.unsqueeze(0).expand(batch_size, -1).to(self.device)
  File "C:\...\__init__.py", line 210, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

从错误堆栈可以看出，问题出现在模型尝试将张量转移到指定设备时，系统检测到当前PyTorch环境没有CUDA支持。

环境配置要点

根据用户反馈，环境配置如下：

• 操作系统：Windows
• Python版本：3.8
• PyTorch版本：1.11.0+cpu（明确为CPU版本）
• 配置文件device参数：已设置为"cpu"
• 模型路径：使用双反斜杠格式

问题根源

经过分析，该问题可能有以下几个原因：

1. PyTorch安装问题：虽然用户安装的是CPU版本的PyTorch(1.11.0+cpu)，但某些依赖可能仍尝试调用CUDA功能。

2. 配置文件加载顺序：可能存在多个配置文件，而实际生效的配置并非用户修改的那个。

3. DeepKE安装问题：项目安装可能不完整或有缓存导致配置未正确更新。

解决方案

根据问题追踪，最终解决方案如下：

1. 重新安装DeepKE：完全卸载后重新安装DeepKE到site-packages目录，确保所有配置和依赖都正确更新。

2. 统一配置文件：确保所有相关配置文件（包括train和predict配置）中的device参数都设置为"cpu"。

3. 路径格式规范：在Windows系统中使用双反斜杠或正斜杠统一路径格式。

4. 环境验证：安装后验证PyTorch确实为CPU版本：

   import torch
   print(torch.__version__)  # 应显示类似1.11.0+cpu
   print(torch.cuda.is_available())  # 应返回False
   

最佳实践建议

对于在CPU环境下使用DeepKE进行NER Few-Shot训练的用户，建议遵循以下步骤：

1. 创建干净的Python虚拟环境
2. 安装CPU版本的PyTorch：

   pip install torch==1.11.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
   

3. 安装DeepKE：

   pip install deepke
   

4. 修改配置文件：
   • 确认所有yaml文件中的device参数设置为"cpu"
   • 检查模型路径格式正确
5. 运行前进行环境验证

总结

DeepKE作为强大的知识抽取工具，在Few-Shot场景下表现优异。遇到CUDA相关问题时，用户应首先确认PyTorch版本与硬件环境匹配，检查所有相关配置文件，并在必要时重新安装项目。通过规范环境配置和安装流程，可以有效避免此类设备不匹配的问题。