Training a YOLOv5 Model on Custom Data

With the images and labels set in their respective folders, we can start training our model! Let's create a data.yaml file that encompasses the information required for the train.py script:

import yaml
config = {'path': os.path.join(os.getcwd(), 'drive-data'),
         'train': 'images/train',
         'test': 'images/test',
         'val': 'images/valid',
         'nc': len(labels),
         'names': labels}
 
with open(os.path.join("drive-data", "data.yaml"), "w") as file:
    yaml.dump(config, file, default_flow_style=False)

Training

We'll initiate the training script as last time - but with a larger number of epochs, and loading the weights from the previous run, rather than from the network that was just pre-trained on MS COCO:

$ python3 yolov5/train.py --batch -1 --epochs 100 --data drive-data/data.yaml --weights yolov5/runs/train/exp/weights/best.pt

This kicks off a relatively short training session, as our dataset is fairly small:

Image sizes 640 train, 640 val
Using 2 dataloader workers
Logging results to yolov5/runs/train/exp2
Starting training for 100 epochs...

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      0/99     7.74G    0.1108   0.02974   0.04983        67       640: 100% 2/2 [00:08<00:00,  4.20s/it]
               Class     Images     Labels          P          R     [email protected] [email protected]:.95: 100% 1/1 [00:00<00:00,  1.15it/s]
                 all         12         17          0          0          0          0
...
Epoch   gpu_mem       box       obj       cls    labels  img_size
     99/99     12.9G   0.02097  0.007557   0.01709        88       640: 100% 2/2 [00:05<00:00,  2.73s/it]
               Class     Images     Labels          P          R     [email protected] [email protected]:.95: 100% 1/1 [00:00<00:00,  4.51it/s]
                 all         12         17      0.735      0.611      0.638      0.413

100 epochs completed in 0.182 hours.
Optimizer stripped from yolov5/runs/train/exp2/weights/last.pt, 14.4MB
Optimizer stripped from yolov5/runs/train/exp2/weights/best.pt, 14.4MB

Validating yolov5/runs/train/exp2/weights/best.pt...
Fusing layers... 
Model summary: 213 layers, 7023610 parameters, 0 gradients, 15.8 GFLOPs
...