YOLO v8 GPU 사용 불가 문제 해결(Could not run ‘torchvision::nms’ with arguments from the ‘CUDA’ backend)

문제

YOLO(You Only Look Once)는 빠르고 정확한 실시간 객체 탐지 알고리즘으로 널리 알려져 있다. YOLO v8은 이러한 시리즈의 최신 버전으로, 더욱 향상된 성능과 기능을 제공한다

그러나 YOLO v8을 사용하면서 GPU를 활용할 때 NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend.와 같은 오류 메시지가 나타날 수 있으며, 이 오류는 CUDA 백엔드에서 ‘torchvision::nms’ 함수를 실행할 때 발생한다.

이 오류의 주요 원인은 CUDA 버전과 torch, torchvision, torchaudio의 버전 간 불일치이며, 라이브러리와 도구의 호환성은 시스템의 성능과 정확도에 중요한 영향을 주고, 버전 불일치는 실행 시 오류를 유발할 수 있다.

해결 방법

라이브러리 및 도구 버전 확인

먼저 현재 설치된 CUDA, torch, torchvision, 및 torchaudio의 버전을 확인합니다.

nvcc --version
pip show torch torchvision torchaudio

PyTorch 웹사이트 방문

PyTorch 공식 웹사이트를 방문한다. 이 사이트에서는 올바른 버전의 torch, torchvision, 및 torchaudio를 설치할 수 있도록 권장 사항이 제시된다.

라이브러리 재설치

웹사이트에서 시스템의 환경에 맞춰 선택하면 필요한 라이브러리와 도구를 재설치할 수 있는 명령어가 나온다.
기존에 설치되었던 torch, torchvision, torchaudio를 삭제하고 Run this Command 설치 문구를 복사하여 재설치한다.

pip uninstall torch torchvision torchaudio
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

라이브러리 재설치 결과 확인

아래 코드를 입력하고 CUDA 버전과 맞게 torch, torchvision, torchaudio가 설치 되었는지 확인한다.

pip show torch torchvision torchaudio

동작 확인하기

CUDA 버전에 맞는 라이브러리들이 설치되었으면, YOLO v8을 실행하여 GPU상에서 정상 동작 되는지 확인한다.

device=0 옵션을 통해 0번 GPU를 선택하였고 67.3ms만에 inferece가 끝난 것을 확인할 수 있다.

위의 단계를 완료하면 정상적으로 YOLO v8이 동작하는 것을 확인할 수 있다.
이제 YOLO v8을 사용하여 GPU에서 빠른 객체 탐지를 즐겨 보자.