MATLAB does not have any function like trainYOLOv4ObjectDetector (for YOLOv4) and trainYOLOXObjectDetector (for YOLOX) to train YOLOv5 models. You'll need to write a custom training loop.
To get started, you need to create a dlnetwork (https://www.mathworks.com/help/deeplearning/ref/dlnetwork.html) object which has the same architecture as YOLOv5.
If you do not want to create a dlnetwork manually, you can try importing a pre-trained YOLOv5 model. The official repository for YOLOv5 has multiple pre-trained networks in different sizes saved in the PyTorch format: https://github.com/ultralytics/yolov5?tab=readme-ov-file#pretrained-checkpoints.
One issue with the models in this repository is that they are not JIT-traced in PyTorch and cannot be loaded using importNetworkFromPyTorch (https://www.mathworks.com/help/deeplearning/ref/importnetworkfrompytorch.html#mw_5b3b9d1e-30ec-435a-a0cc-9dd86bbb76c8). Also, the models use the float16 data type, which is not supported by MATLAB.
To get arond these issues, you can use the following workflow:
- Install Python and PyTorch.
- Clone the offical git repository to a folder of your choice using the command git clone https://github.com/ultralytics/yolov5.git.
- Download any of the pre-trained checkpoints based on your preferences. I am using the smallest yolov5n.pt for this example.
- Write and execute a Python script with the following code:
import torch
# Load the model
# TODO: Change the filename based on what you downloaded
model = torch.load("yolov5n.pt")
model = model["model"]
# Convert the model to float32 since MATLAB does not support float16
model = model.float()
# Create a dummy image to run through the model
# TODO: Change this to the size of your images
torch_input = torch.randn(1, 3, 224, 224)
# Export to ONNX format
onnx_program = torch.onnx.export(model, (torch_input,), "yolov5n.onnx")
The script will save the model in the ONNX format, which can then be loaded in MATLAB using the importNetworkFromONNX function.
net = importNetworkFromONNX("yolov5n.onnx");
Refer to the following example which shows how to create a custom training loop: https://www.mathworks.com/help/deeplearning/ug/train-network-using-custom-training-loop.html
Refer to the following resources for more information:
- Pretrained networks from external platforms - https://www.mathworks.com/help/deeplearning/networks-from-external-platforms.html
- importNetworkFromONNX documentation - https://www.mathworks.com/help/deeplearning/ref/importnetworkfromonnx.html
- Exporting models to ONNX in PyTorch - https://pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html.
- trainYOLOXObjectDetector documentation - https://www.mathworks.com/help/vision/ref/trainyoloxobjectdetector.html
Hope this helps!
