Yolov7训练自定义数据集遇torch.cuda.amp导入错误求助

环境信息 #

• CUDA版本：9.0
• Python版本：3.7.1
• PyTorch版本：1.1.0
• TorchVision版本：0.3.0
• 执行的训练命令：

(venv37) PS D:\aiml_python_code\yolov7-main> python .\train.py --batch-size 8 --epochs 100 --img 640 640 --data .\data\car_dataset.yaml --hyp .\data\hyp.scratch.custom.yaml --cfg .\cfg\training\yolov7_custom.yaml --name yolov7-custom --weights yolov7.pt --device 0 --workers 1

报错详情 #

Traceback (most recent call last):
File ".\train.py", line 19, in
from torch.cuda import amp
ImportError: cannot import name 'amp' from 'torch.cuda' (D:\aiml_python_code\yolov7-main\venv37\lib\site-packages\torch\cuda_init_.py)

注：torch.cuda.is_available()返回True，CUDA环境正常，问题核心是torch1.1.0未内置torch.cuda.amp模块（该模块在torch1.6.0及以上版本才引入），而CUDA9.0限制无法升级PyTorch版本。

兼容解决方案 #

方案1：添加参数禁用混合精度训练 #

在训练命令末尾追加--no-half参数，强制关闭混合精度训练逻辑，代码将不再调用amp模块：

python .\train.py --batch-size 8 --epochs 100 --img 640 640 --data .\data\car_dataset.yaml --hyp .\data\hyp.scratch.custom.yaml --cfg .\cfg\training\yolov7_custom.yaml --name yolov7-custom --weights yolov7.pt --device 0 --workers 1 --no-half

方案2：修改train.py移除amp依赖 #

如果方案1不生效，可手动调整代码：

1. 打开train.py，注释或删除第19行的from torch.cuda import amp
2. 找到并删除amp.GradScaler()的初始化代码
3. 替换所有with amp.autocast():代码块为普通前向传播逻辑
4. 将scaler.scale(loss).backward()改为loss.backward()，移除梯度缩放逻辑

优先推荐方案1，无需修改代码，操作更简便。

内容的提问来源于stack exchange，提问作者Mogli141