YOLOv8 torch 오류 해결 - Could not run ‘torchvision::nms’ with arguments from the ‘CUDA’ backend

YOLOv8을 실행 중 다음과 같은 문제가 발생하였다.

NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend

NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend 오류는 PyTorch와 torchvision을 사용하여 객체 검출(Object Detection) 작업을 할 때 발생할 수 있다.

이 오류는 일반적으로 NVIDIA CUDA를 지원하는 GPU에서 발생하는데, 특히 torchvision의 nms (Non-Maximum Suppression) 함수를 호출할 때 CUDA 백엔드가 제대로 동작하지 않을 때 발생한다.

해당 오류가 발생한다면 다음과 같은 절차를 따라 해결하면 된다.

 

1. CUDA Toolkit 설치 

우선 CUDA ToolKit을 설치한다. CUDA ToolKit를 이미  설치 한 경우 이 과정은 건너 뛰어도 된다. 

CUDA Toolkit Archive

이곳에서 적절한 CUDA버전을 설치한다.

자신의 GPU에 맞는 버전이 뭔지 모르겠다면 아래를 참고해서 다운 받으면 된다.

자신의 GPU에 맞는 CUDA 버전 설치하기

 

CUDA 가 완료된다면 다음 명령어를 실행해서 CUDA가 제대로 설치되었는지 확인한다. 

nvidia-smi
nvcc --version

 

반응형

 

2. torch torchvision torchaudio 설치 및 버전 일치

다음은 torch torchvision torchaudio의 버전을 확인한다.

(1) Python 코드로 확인

import torch
import torchvision
import torchaudio

print("PyTorch version:", torch.__version__)
print("torchvision version:", torchvision.__version__)
print("torchaudio version:", torchaudio.__version__)

혹은

(2) CLI로 확인

pip show torch torchvision torchaudio

 

각 버전 뒤에 CUX.X.X 의 버전이 같아야한다. 만약 잘못된 버전이나 버전 뒤에 +CUX.X.X가 없다면 torch torchvision torchaudio을 재설치 해야한다.

 

 

3. torch torchvision torchaudio 설치 및 재설치

(1) torch torchvision torchaudio 제거

우선 아래 명령어를 통해 현재 설치된 torch torchvision torchaudio 라이브러리를 찾는다.

pip list

위와 같이 모두 있는경우 모두 제거해준다.

pip uninstall torch
pip uninstall torchaudio
pip uninstall torchvision

 

 

이제 Pytouch 공식 홈페이지에 가서 맞는 정식 버전을 설치한다.

Start Locally

나는 Window환경에 Python으로 개발하고 있으며 CUDA는 21.1버전이 호환됨으로 다음과 같이 설청하였다. 만약 아나 콘다를 쓰고 있거나 CUDA 11.8버전을 사용해야한다면 그에 맞춰서 설치하면 된다.

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

 

잘 실행했다면 다시 pip show torch torchvision torchaudio을 실행해서 버전을 확인한다. 만약 모두 버전 뒤에 + CUX.X.X가 붙어있다면 제대로 설치 된 것이다.