Can i use custom model as pretrained model, and train again?

Search before asking

[X] I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

Hello,

I have two questions:

I used yolov5s.pt as pretrained model when i firstly train my custom model, however, my custom model is really special, it's not about classifying human, cars or trees. It's a defect detection, which is used to detect one object defective or not. In this case, it seems the pretrained model yolov5s.pt does not has too much connection with my custom dataset. In this case, is it still good to use yolov5s.pt as pretrained model?
Now i have already trained my custom model, and want to improve the perfromance. Can i use the current custom model as pretrained model and train it again (use new dataset), does this method works?

Additional

No response

👋 Hello @shengjieH, 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 a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Requirements

Python>=3.8.0 with all requirements.txt installed including PyTorch>=1.8. 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):

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, validation, inference, export and benchmarks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

Introducing YOLOv8 🚀

We're excited to announce the launch of our latest state-of-the-art (SOTA) object detection model for 2023 - YOLOv8 🚀!

Designed to be fast, accurate, and easy to use, YOLOv8 is an ideal choice for a wide range of object detection, image segmentation and image classification tasks. With YOLOv8, you'll be able to quickly and accurately detect objects in real-time, streamline your workflows, and achieve new levels of accuracy in your projects.

Check out our YOLOv8 Docs for details and get started with:

pip install ultralytics

@shengjieH hello,

Thank you for your questions! Let's address each one:

Using yolov5s.pt as a Pretrained Model for Defect Detection: While yolov5s.pt is trained on the COCO dataset, which includes a wide variety of objects, it can still be beneficial as a starting point for your defect detection task. The pretrained weights can help the model learn general features such as edges, textures, and shapes, which can be useful even if your specific task is different. However, if you find that the pretrained model does not significantly improve your results, you might consider training from scratch or using a smaller, more relevant dataset to pretrain your model.
Using a Custom Model as a Pretrained Model for Further Training: Yes, you can definitely use your already trained custom model as a pretrained model for further training. This approach is known as fine-tuning and can be particularly effective if you have a new dataset that is similar to your original dataset but with some variations. To do this, you can simply load your custom model's weights and continue training with the new dataset. Here’s a quick example of how you can do this:
```
import torch
from yolov5 import train

# Load your custom model
model = torch.load('path/to/your/custom_model.pt')

# Continue training with the new dataset
train.run(data='path/to/your/new_dataset.yaml', weights='path/to/your/custom_model.pt', epochs=50)
```
This method allows the model to leverage the knowledge it has already gained from the initial training and adapt to the new data, potentially improving performance.

For more detailed guidance, you can refer to our Train Custom Data tutorial.

I hope this helps! If you have any further questions, feel free to ask. 😊

@glenn-jocher Thank you so much for the quick the detailed reply!!

Actually I met a problem during my training. When detecting small objects using yolov5, the light condition is a big issue which can cause light reflection and a serial of questions, especailly when the camera is not really good and has low resolution. Is there any methods that can reduce the light influence? For example, by using some data preprocessing technology or controling hyp parameters?

Hello @shengjieH,

Thank you for your kind words! Addressing lighting conditions and reflections can indeed be challenging, especially with small objects and lower-resolution cameras. Here are a few strategies you can employ to mitigate these issues:

Data Preprocessing Techniques

Data Augmentation: Use data augmentation techniques to simulate various lighting conditions during training. This can help the model become more robust to changes in lighting. You can apply transformations such as brightness adjustment, contrast adjustment, and random shadows. Here’s an example using Albumentations:
```
import albumentations as A

transform = A.Compose([
   A.RandomBrightnessContrast(p=0.5),
   A.RandomShadow(p=0.5),
   A.HueSaturationValue(p=0.5)
])
```

Histogram Equalization: Apply histogram equalization to improve the contrast of images. This can help in highlighting features that might be washed out due to lighting conditions.

import cv2

def equalize_histogram(image):
   img_yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
   img_yuv[:, :, 0] = cv2.equalizeHist(img_yuv[:, :, 0])
   return cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)

Hyperparameter Tuning

Learning Rate: Adjust the learning rate to ensure the model learns effectively from the augmented data. Sometimes, a lower learning rate can help the model adapt better to variations in the data.
Anchor Boxes: Ensure that the anchor boxes are well-suited for the size of the small objects you are detecting. You can use the kmeans algorithm to optimize anchor boxes for your dataset.
```
from yolov5.utils.autoanchor import kmean_anchors

kmean_anchors('path/to/your/dataset.yaml')
```

Additional Tips

Use a Better Camera: If possible, upgrading to a higher-resolution camera can significantly improve detection performance by providing more detailed images.
Lighting Control: Try to control the lighting environment where the images are captured. Consistent lighting can reduce the variability that the model needs to handle.
Post-Processing: Implement post-processing techniques such as non-maximum suppression (NMS) with appropriate thresholds to filter out false positives that might be caused by reflections.

For more detailed guidance on these techniques, you can explore our Tips for Best Training Results guide.

I hope these suggestions help improve your model's performance under varying lighting conditions. If you have any further questions or need additional assistance, feel free to ask. 😊

ultralytics / yolov5