Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

model.val() may can not work on converted onnx file. #451

Open
JIAOJINYU opened this issue Sep 27, 2024 · 0 comments
Open

model.val() may can not work on converted onnx file. #451

JIAOJINYU opened this issue Sep 27, 2024 · 0 comments

Comments

@JIAOJINYU
Copy link

The pretrained .pt files (I get it through wget https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10{n/s/m/b/l/x}.pt) can be evaulated on coco with

from ultralytics import YOLOv10
model = YOLOv10('yolov10n.pt')
model.val(data='coco.yaml', batch=8) 

After I converted the .pt file to onnx file, it can not be evaluated on coco val with

from ultralytics import YOLOv10
model = YOLOv10('yolov10n.onnx')
model.val(data='coco.yaml') 

The error is

Traceback (most recent call last):
  File "/workspace/yolov10_acc/val.py", line 5, in <module>
    model.val(data='coco.yaml')  
  File "/workspace/yolov10_acc/yolov10/ultralytics/engine/model.py", line 516, in val
    validator(model=self.model)
  File "/root/.pyenv/versions/yolov10/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/workspace/yolov10_acc/yolov10/ultralytics/engine/validator.py", line 187, in __call__
    preds = self.postprocess(preds)
  File "/workspace/yolov10_acc/yolov10/ultralytics/models/yolov10/val.py", line 18, in postprocess
    boxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, self.nc)
  File "/workspace/yolov10_acc/yolov10/ultralytics/utils/ops.py", line 852, in v10postprocess
    assert(4 + nc == preds.shape[-1])
AssertionError"

Since the onnx's shape is [1, 300, 6] as following picture.
image
I think the output post-processed detections that have already undergone the NMS as a result it is not appropriate for the evaluation process in the YOLOv10DetectionValidator. Evaluation process typically expect the model to output raw predictions.(may include bbox coordinates and class probabilities for all classes not just the highest-scoring class)

For the onnx's output shape, In my understanding, it is in the shape [batch_size, num_predict, 6], where 6 corresponds to [x1, y1, x2, y2, score, label]. To solve this unaligned shape and postprocessing, I adjust the postprocess method in the YOLOv10DetectionValidator class to correctly handle the shape of preds from the .pt model and converted onnx file.
I uploaded the modified yolov10/ultralytics/models/yolov10/val.py and added the necessary comments to the modified code. Note that I only experimented with yolov10n.pt and yolov10n.onnx. For other models there is no guarantee that it will work.

from ultralytics.models.yolo.detect import DetectionValidator
from ultralytics.utils import LOGGER, ops
import torch
import numpy as np

class YOLOv10DetectionValidator(DetectionValidator):
    """
    Custom validator for YOLOv10 that handles evaluation for both .pt and .onnx models.
    """
    def __init__(self, *args, **kwargs):
        """
        Initializes the validator by setting up necessary attributes.
        """
        super().__init__(*args, **kwargs)
        # Ensure that JSON saving is enabled for COCO dataset
        self.args.save_json |= self.is_coco
        self.nc = None  # Number of classes will be determined later

    def __call__(self, trainer=None, model=None):
        """
        Overrides the __call__ method to set up the number of classes (nc) before validation.
        """
        self.model = model  # Store the model
        if self.model is None:
            raise ValueError('Model is None in validator call.')

        # Determine the number of classes (nc)
        try:
            if hasattr(self.model, 'model') and isinstance(self.model.model, (list, tuple)):
                # For .pt models: access the last layer to get nc
                last_layer = self.model.model[-1]
                self.nc = getattr(last_layer, 'nc', getattr(last_layer, 'num_classes', None))
            elif hasattr(self.model, 'names'):
                # For models with 'names' attribute: get nc from length of names
                self.nc = len(self.model.names)
            elif self.data and 'nc' in self.data:
                # Get nc from data configuration
                self.nc = self.data['nc']
            else:
                # Default to 80 classes (e.g., for COCO dataset)
                self.nc = 80
                LOGGER.warning('Number of classes not found. Assuming nc=80.')
            if self.nc is None:
                raise ValueError('Number of classes (nc) could not be determined from the model or data.')
        except Exception as e:
            raise ValueError(f'Error determining number of classes: {e}')

        # Proceed with validation using the parent class's __call__ method
        super().__call__(trainer=trainer, model=model)

    def postprocess(self, preds):
        """
        Postprocesses predictions to handle outputs from both .pt and .onnx models.

        Args:
            preds: The raw predictions from the model.

        Returns:
            Processed predictions ready for evaluation.
        """
        # Handle different types of predictions and convert to torch.Tensor if necessary
        if isinstance(preds, dict):
            preds = preds.get("one2one", preds)

        if isinstance(preds, (list, tuple)):
            preds = preds[0]  # Take the first element if preds is a list or tuple

        if isinstance(preds, np.ndarray):
            preds = torch.from_numpy(preds)  # Convert numpy array to torch.Tensor

        if not isinstance(preds, torch.Tensor):
            raise TypeError(f'Predictions should be a torch.Tensor but got {type(preds)}.')

        # Process predictions based on their dimensions and shapes
        if preds.ndim == 3:
            if preds.shape[1] == self.nc + 4:
                # Raw outputs from .pt model with shape [batch_size, nc + 4, num_predictions]
                preds = preds.permute(0, 2, 1)  # Transpose to [batch_size, num_predictions, nc + 4]

                # Postprocess predictions (e.g., apply NMS)
                try:
                    boxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, self.nc)
                except Exception as e:
                    raise RuntimeError(f'Error in v10postprocess: {e}')
                bboxes = ops.xywh2xyxy(boxes)  # Convert boxes to [x1, y1, x2, y2] format
                predn = torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)
                return predn  # Return tensor of predictions

            elif preds.shape[-1] == 6:
                # Processed outputs from .onnx model with shape [batch_size, num_predictions, 6]
                # Remove batch dimension if necessary
                if preds.shape[0] == 1:
                    preds = preds[0]
                return [preds]  # Return list of predictions

            else:
                raise ValueError(f'Unknown prediction format for 3D tensor with shape {preds.shape}.')

        elif preds.ndim == 2 and preds.shape[-1] == 6:
            # Processed outputs from .onnx model with shape [num_predictions, 6]
            return [preds]  # Return list of predictions

        else:
[onnx_and_val.zip](https://github.com/user-attachments/files/17162353/onnx_and_val.zip)

            raise ValueError(f'Unknown prediction format with tensor shape {preds.shape}.')

I know that this is only a temporary solution, if you have a better solution please share it with me.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant