The recognition result of classify/predict.py is very poor?

[small-code-cat]

Search before asking

• I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

I use classify/train.py to train the target classification, the model I have used yolov5s-cls.pt and efficient_v2_s, the first model was trained 50 times and the second one was trained 10 times. The loss has been decreasing and the accuracy is also good; however, the results of classify/predict.py are very poor
The prediction results are as follows.

Additional

No response

[github-actions]

👋 Hello @small-code-cat, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

If this is a 🐛 Bug Report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online W&B logging if available.

For business inquiries or professional support requests please visit https://ultralytics.com or email support@ultralytics.com.

Requirements

Python>=3.7.0 with all requirements.txt installed including PyTorch>=1.7. To get started:

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab and Kaggle notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), validation (val.py), inference (detect.py) and export (export.py) on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[small-code-cat]

The dataset I use is minist. I hope someone can help me

[glenn-jocher]

@small-code-cat MNIST training runs correctly for me. I see 97.5% top 1 accuracy after two epochs with all default settings.

[small-code-cat]

Yes, my training results are as good as you, but look at my prediction results are very poor. I posted a screenshot of the prediction on it.

[glenn-jocher]

@small-code-cat training is at imgsz 128, so you might want to predict at the same size.

[small-code-cat]

I don't quite understand. I didn't specify imgsz during training. Why do I need to specify when predicting? Isn't this automatically processed by the framework?

[glenn-jocher]

@small-code-cat see train.py and predict.py argparsers for details, i.e.:

yolov5/classify/predict.py

Lines 183 to 206 in 4d50cd3

	def parse_opt():
	parser = argparse.ArgumentParser()
	parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s-cls.pt', help='model path(s)')
	parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam')
	parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
	parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[224], help='inference size h,w')
	parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
	parser.add_argument('--view-img', action='store_true', help='show results')
	parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
	parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
	parser.add_argument('--augment', action='store_true', help='augmented inference')
	parser.add_argument('--visualize', action='store_true', help='visualize features')
	parser.add_argument('--update', action='store_true', help='update all models')
	parser.add_argument('--project', default=ROOT / 'runs/predict-cls', help='save results to project/name')
	parser.add_argument('--name', default='exp', help='save results to project/name')
	parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
	parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
	parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
	parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
	opt = parser.parse_args()
	opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
	print_args(vars(opt))
	return opt

[small-code-cat]

Thank you very much, it helped me a lot, it turns out that the imgsz of the two codes are not the same

[small-code-cat]

By the way, there is another question I want to ask. When using torch.hub.load to load the local classification model, it will report an error when going to predict images.

[glenn-jocher]

@small-code-cat ClassificationModel and SegmentationModel types are not yet fully supported by AutoShape, but you can still load them with PyTorch Hub and pass them torch tensors. Note you must apply your own pre and post processing:

[glenn-jocher]

@small-code-cat good news 😃! Your original issue may now be partially resolved ✅ in PR #9469. This PR aligns all default Classification task image sizes at 224 pixels. This should help avoid confusion in the future.

Note that many datasets like MNIST with very small image sizes (i.e. 28x28 pixels) can achieve good results at lower image resolutions like --imgsz 128, but ImageNet default is 224.

To receive this update:

• Git – git pull from within your yolov5/ directory or git clone https://github.com/ultralytics/yolov5 again
• PyTorch Hub – Force-reload model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
• Notebooks – View updated notebooks
• Docker – sudo docker pull ultralytics/yolov5:latest to update your image

Thank you for spotting this issue and informing us of the problem. Please let us know if this update resolves the issue for you, and feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀!

[small-code-cat]

Thank you very much, this update is very useful, I hope yolov5 can do better and better, cheers!

[small-code-cat]

@small-code-cat ClassificationModel and SegmentationModel types are not yet fully supported by AutoShape, but you can still load them with PyTorch Hub and pass them torch tensors. Note you must apply your own pre and post processing:

I tried it and it's running, thanks a lot!

[small-code-cat]

There is another question. After I train the model with gpu on Linux, loading the model on macos will report an error. What's the reason?

[glenn-jocher]

@small-code-cat perhaps your paths are incorrect. I would just load as in the default example from ultralytics/yolov5

[small-code-cat]

I have tried this, and I will report this error.

[small-code-cat]

@small-code-cat perhaps your paths are incorrect. I would just load as in the default example from ultralytics/yolov5

So is there any way to solve this problem?

[small-code-cat]

@small-code-cat perhaps your paths are incorrect. I would just load as in the default example from ultralytics/yolov5

The problem has been solved, because I have a python file called utils in the same level directory, change the name and it is solved

[small-code-cat]

How do you remove the last few prediction layers when loading the model？ I want to use some of the layers to do kmeans clustering

[glenn-jocher]

@small-code-cat you can access the model's layers through the model's model.model attribute, which contains the actual YOLOv5 PyTorch model, and then modify it as needed. For example:

import torch
from models.yolo import Model

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Access the actual YOLOv5 PyTorch model
yolo_model = model.model

# Remove the last few prediction layers
yolo_model.model[-1].pop()  # Remove the last layer from the YOLO layer list
yolo_model.model[-1].pop()  # Remove the second-to-last layer from the YOLO layer list

# Perform kmeans clustering using the modified model
# ...

Please note that modifying the model in this way will alter its behavior, so it's important to understand the implications of removing specific layers.

For more details on working with the YOLOv5 PyTorch model, you can refer to the Ultralytics YOLOv5 documentation.