YOLOv8でのエラー(Could not run 'torchvision::nms' with arguments from the 'CUDA' backend.)を解決

状況

公式サイトの通りPyTorch, Torchvision, Torchaudioをインストール

YOLOv8モデルを実行すると

from ultralytics import YOLO

dataset_path = "datasets"

model = YOLO("yolov8n-cls.pt")

results = model.train(data=dataset_path, epochs=100, device="cuda")

エラーが発生

~略~
Could not run 'torchvision::nms' with arguments from the 'CUDA' backend.
~略~

解決

pip listでパッケージを確認

torch                        2.2.0+cu118
torchaudio                   2.2.0+cu118
torchvision                  0.17.0

torchvisionがgpu verになっていない

gpu verのtorchvisionをインストール

pip install torchvision==0.17.0+cu118 -f https://download.pytorch.org/whl/torch_stable.html

pip listで正常にインストールされていることを確認

torch                        2.2.0+cu118
torchaudio                   2.2.0+cu118
torchvision                  0.17.0+cu118

再度YOLOv8モデルを実行すると別のエラーが発生

RuntimeError:p
        An attempt has been made to start a new process before the
        current process has finished its bootstrapping phase.

        This probably means that you are not using fork to start your
        child processes and you have forgotten to use the proper idiom
        in the main module:

            if __name__ == '__main__':
                freeze_support()
                ...

        The "freeze_support()" line can be omitted if the program
        is not going to be frozen to produce an executable.

プログラムを修正

from ultralytics import YOLO

dataset_path = "datasets"

if __name__ == "__main__":
    model = YOLO("yolov8n-cls.pt")

    results = model.train(data=dataset_path, epochs=100, imgsz=200, device="cuda")

無事正常に動作

Starting training for 100 epochs...

      Epoch    GPU_mem       loss  Instances       Size
      1/100     0.539G      1.551          6        224: 100%|██████████| 79/79 [00:05<00:00, 15.60it/s]
               classes   top1_acc   top5_acc: 100%|██████████| 9/9 [00:00<00:00, 16.66it/s]
                   all      0.494          1