If you’re working on deep learning projects using YOLOv8 (You Only Look Once), you may have encountered the error: “CUDA error ‘torchvision::nms’ With YOLOv8 Model on NVIDIA GTX Card”. This error typically arises in environments utilizing NVDIA’s CUDA for GPU(graphics processing units) and video cards acceleration. This blog post provides a how-to guide with a step-by-step explanation to identify and resolve this error, allowing you to focus and continue building high-performance machine-learning models.

What Does the Error Mean?

We are using NVIDIA GeForce GTX 1050 GPU with 4GB VRAM to train a YOLOv8 model. NVDIA’s GTX cards being economical solution when you are training models locally. There are high performing GPU cards like the latest one’s RTX series (ray tracing). Also one can utilize the GPU cloud solutions from Google cloud or Amazon cloud to train the model. Upon executing the training command, an unexpected error surfaced:

NotImplementedError: Could not run ‘torchvision::nms’ with arguments from the ‘CUDA’ backend.

specifically relates to the torchvision package’s Non-Maximum Suppression (NMS) operation. The error usually means that the nms function is attempting to utilize GPU acceleration but is unable to due to incompatible versions.

Why the Error Occurs?

– CUDA Version Installed: 12.2
– PyTorch Version: 2.3.1+cu118 (Built for CUDA 11.8)
– Error Context: Mismatch between the installed CUDA version and the CUDA version PyTorch was built against, leading to incompatibility issues with `torchvision`.

To avoid these issues, always check compatibility requirements on the official PyTorch installation page.

Step-by-Step Troubleshooting

Step 1: Verifying CUDA Installation

The first step was to confirm the CUDA installation and ensure that the system recognized the GPU:

nvcc --version

If NVDIA’s CUDA compiler driver is installed, This should produce the output along with driver details. If not it will prompt you to install. The output we are interested in is driver version and it should look something like

CUDA compilation tools, release 12.2, V12.2.140

Note that, release version may vary based on your ‘nvcc’ installation, once can use ‘nvidia-smi’ command as well.

Step 2: Setting Up a Python Virtual Environment

To manage dependencies effectively and avoid conflicts with current systems installation, It is recommended to use a virtual environment. Virtual environment is created using `virtualenv` and follows below shell commands

pip install virtualenv
virtualenv yolov8_env

Activate the environment, On Windows:

yolov8_env\Scripts\activate

Step 3: Installing PyTorch with Compatible CUDA Support

Given the CUDA version discrepancy, two options were considered:

Downgrading CUDA to a Supported Version (Recommended), Action: Uninstall CUDA 12.2 and install CUDA 11.7.

Attempting to Install PyTorch with CUDA 12.2 (If Supported),  Action: Check PyTorch’s compatibility with CUDA 12.2 and install accordingly.

Opting for the first approach ensured compatibility, aligning PyTorch with CUDA 11.8.

Step 4: Installing Ultralytics YOLOv8

With PyTorch set up, the next step was to install the YOLOv8 package:

pip install ultralytics

Step 5: Preparing the Dataset

The dataset was organized following YOLOv8’s expected structure, including image and label directories, and a `data.yaml` configuration file specifying paths, number of classes (`nc`), and class names.

Step 6: Optimizing for Limited GPU Memory

Given the GTX 1050’s 4GB VRAM, several optimizations were implemented:

– Reducing Batch Size: From an initial `batch=4` to smaller sizes like `batch=2`.
– Decreasing Image Size: Lowering `imgsz` from 416 to 320 or 256 to reduce memory usage.
– Enabling Mixed Precision: Utilizing `–fp16` to leverage Automatic Mixed Precision, cutting memory consumption.
– Implementing Gradient Accumulation: Using `–accumulate 2` to simulate larger batch sizes without increasing actual batch size.

yolo train data=data.yaml model=yolov8s.pt epochs=100 imgsz=416 batch=16 device=0

Step 7: Ensuring Torchvision Compatibility

To resolve the initial `torchvision::nms` CUDA error, the following steps were taken:

Verifying package Installation

import torchimport 
import torchvision
print("PyTorch Version:", torch.__version__)
print("CUDA Available:", torch.cuda.is_available())
print("CUDA Version:", torch.version.cuda)
print("torchvision Version:", torchvision.__version__)

Reinstalling torchvision with Matching CUDA Support

pip uninstall torchvision -y
pip install torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118

Testing NMS Operation

from torchvision.ops import nms
boxes = torch.tensor([[10, 10, 20, 20], [15, 15, 25, 25]], dtype=torch.float).cuda()
scores = torch.tensor([0.9, 0.75], dtype=torch.float).cuda()
selected_indices = nms(boxes, scores, iou_threshold=0.5)
print("NMS Selected Indices:", selected_indices)

These steps ensured that `torchvision` was correctly installed and compatible with the PyTorch version in use, resolving the initial CUDA backend issue.

Step 8: Monitoring GPU Memory Usage

Throughout the training process, GPU memory utilization was monitored using `nvidia-smi` to ensure that optimizations were effective and that the GPU resources were being utilized efficiently.

Achieving Success: What Worked

Aligning PyTorch and Torchvision with Compatible CUDA Versions

Ensuring that both PyTorch and `torchvision` were installed with matching CUDA support (`cu118`) eliminated compatibility issues. Reinstalling `torchvision` with the specific CUDA version addressed the initial NMS error, allowing the model to leverage GPU acceleration effectively.

Optimizing Training Parameters for Limited VRAM

By adjusting key training parameters, the GPU’s limited memory was utilized efficiently:

– Batch Size: Reducing the batch size to 2 prevented memory overloads.
– Image Size: Keeping the image size at 416 balanced between performance and memory usage.
– Mixed Precision: Enabling `–fp16` allowed for reduced memory consumption without significantly compromising model accuracy.
– Gradient Accumulation: Using `–accumulate 2` simulated a larger batch size, enhancing training stability and performance.

Verifying and Updating Dependencies

Regularly updating Ultralytics YOLOv8 and ensuring that all dependencies were correctly installed and compatible played a crucial role in maintaining a stable training environment.

Final Training Command That Worked for us

Combining all successful adjustments, the final training command that facilitated efficient training within the GPU’s memory constraints was:

yolo train data=data.yaml model=yolov8s.pt epochs=100 imgsz=416 batch=2 device=0

Key Parameters:

– `data=data.yaml`: Specifies the dataset configuration.
– `model=yolov8s.pt`: Uses the correct YOLOv8 small pretrained model.
– `epochs=100`: Sets the number of training epochs.
– `imgsz=416`: Maintains a balanced image size for performance and memory usage.
– `batch=2`: Keeps the batch size low to fit within 4GB VRAM.

Monitoring and Validation – Post-correction, the training process showed consistent GPU memory usage around 0.9 GB, indicating that optimizations were effective. The initial loss metrics were within expected ranges, and the training proceeded without further CUDA-related interruptions.

Conclusion

The error “CUDA error ‘torchvision::nms’ With YOLOv8 Model on NVIDIA GTX Card” often results from mismatched versions of PyTorch, CUDA, and torchvision. By following the steps outlined here, you can troubleshoot and resolve the issue. Training a YOLOv8 model on an NVIDIA GTX 1050 with 4GB VRAM is entirely achievable with careful configuration and optimization. This journey underscores the importance of meticulous troubleshooting, aligning software dependencies, and adapting training parameters to suit hardware capabilities.