在使用WongKinYiu/PyTorch_YOLOv4进行训练时,报错:
Traceback (most recent call last):
File "train.py", line 15, in <module>
from torch.cuda import amp
ImportError: cannot import name 'amp'
1.只有PyTorch1.6版本以上才可以从torch.cuda中import amp;
2.否则需要自行安装apex,源代码改为
from apex import amp
1.检查PyTorch和CUDA版本是否对应;
https://pytorch.org/get-started/previous-versions/
2.更新PyTorch版本为1.6及以上(建议,因为模型中可能还有其他低版本不兼容的代码)
或
安装apex
git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py install --cpp_ext --cuda_ext
(有时会安装失败,检查PyTorch和CUDA版本 或者 去掉--cuda_ext 便可顺利安装)
并将from torch.cuda import amp
改为from apex import amp
。