pytorch yolov5 Training custom data

---

---

Preface

Environmental Science

python: 3.9.7

torch: 1.10.2

labelimg: 1.8.6

#yolov5 https://github.com/ultralytics/yolov5
#pytorch https://pytorch.org/
#labelimg https://github.com/tzutalin/labelImg

paddleocr There are three models det testing cls Direction rec distinguish

---

One 、 Create an environment

1. install miniconda

   https://blog.csdn.net/mtl1994/article/details/114968140
   

2. Create an environment

   #linux  It needs to be done first  source
   conda create -n pytorch_yolov5 python=3.9.7  --channel https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
   

3. Enter the environment after installation

   conda activate pytorch_yolov5
   

Two 、 Installation environment

1. pytorch

   # Select corresponding cuda/cpu edition 
   pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
   

2. yolov5

   # Download source code 
   git clone https://github.com/ultralytics/yolov5  # clone
   cd yolov5
   # Installation dependency 
   python -m pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/
   

3. Test it

   import torch
   
   # Model
   model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5n - yolov5x6, custom
   
   # Images
   img = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list
   
   # Inference
   results = model(img)
   
   # Results
   results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
   

4. Choose a model

   https://github.com/ultralytics/yolov5/releases

   Model parameters
   

   The model I use is

   yolov5x
   

3、 ... and 、 Mark the picture

I use labelimg

# install 
pip install labelimg
# open 
labelimg

After marking, there will be two directories 、 A save picture 、 A deposit txt

Four 、 Training

1. Segmentation data

def make_datasets(txt_path, img_path, out="./out", split_rate=0.05):
    txt_dir = Path(txt_path)
    img_dir = Path(img_path)
    out_dir = Path(out)
    dataset = []

    train_label = out_dir / "labels/train2007/"
    train_image = out_dir / "images/train2007/"
    test_label = out_dir / "labels/test2007/"
    test_image = out_dir / "images/test2007/"
    train_label.mkdir(parents=True, exist_ok=True)
    train_image.mkdir(parents=True, exist_ok=True)
    test_label.mkdir(parents=True, exist_ok=True)
    test_image.mkdir(parents=True, exist_ok=True)

    """
     Filter empty txt
    """
    for item in txt_dir.rglob("*.txt"):
        if item.read_text() != "":
            dataset.append(item)
    """
     Split training set , Verification set 
    """
    tv = random.sample(dataset, int(len(dataset) * split_rate))
    """
     Assembly data 
    """

    print(len(dataset))
    for item in dataset:
        for jpg in  img_dir.rglob(item.stem + ".jpg"):
            if item in tv:
                print(jpg, test_image / jpg.name)
                # jpg.replace(test_image / jpg.name)
                shutil.copy(str(jpg), test_image / jpg.name)
                shutil.copy(str(item), test_label / item.name)

                # item.replace(test_label / item.name)
            else:
                print(jpg, train_image / jpg.name)
                # jpg.replace(train_image / jpg.name)
                # item.replace(train_label / item.name)
                shutil.copy(str(jpg), train_image / jpg.name)
                shutil.copy(str(item), train_label / item.name)

After execution , Directory structure

2. Modify the training model yml

3. Start training

nohup python train.py --img 640 --batch 32 --epochs 600 --data voc.yaml --weights yolov5s.pt --device 0,1,2,3 &

The training results are saved in runs/train/ In the incremental running directory , for example runs/train/exp2,runs/train/exp3

5、 ... and 、 Trained models , With new data, we need to continue training

1. Use transfer learning

Designated during training weights by The output of the last training

6、 ... and 、 Common commands

1. Training

nohup python train.py --img 640 --batch 32 --epochs 600 --data wp_voc.yaml --weights runs/train/exp27/weights/best.pt --device 0,1,2,3 &

2. distinguish

python detect.py --weights runs/train/exp6/weights/best.pt --source ../datasets/infer/2022-2-24/

3. export onnx

python export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1

summary