在YoloV5-Lite目标检测之“安装推理”中，我们完成了安装和预训练权重的推理，下面介绍自定义训练数据、模型转换(ncnn)

1 训练数据准备

.
├── train
│   ├── 000000000049.jpg
│   ├── 000000000049.txt
......
│   ├── 000000581880.txt
│   ├── 000000581900.jpg
│   └── 000000581900.txt
└── val
    ├── 000000000139.jpg
    ├── 000000000139.txt
......
├── 000000581357.jpg
    └── 000000581357.txt

3 directories, 46201 files

2 配置文件修改

复制要训练的yml配置

cd data
cp ./data/coco128.yaml custom.yaml

修改配置，主要是目录

# path
train: coco_2017/train
val: coco_2017/val

# number of classes
nc: 1

# class names
names: [ 'person']

复制模型配置

cp ./models/./models/v5Lite-s.yaml ./v5Lite-s-custom.yaml

修改，主要是classes

nc: 1  # number of classes

3 代码修改

有一些版本兼容bug修一下

diff --git a/train.py b/train.py
index 787c243..ff3096e 100644
--- a/train.py
+++ b/train.py
@@ -302,7 +302,7 @@ def train(hyp, opt, device, tb_writer=None):
             # Forward
             with amp.autocast(enabled=cuda):
                 pred = model(imgs)  # forward
-                loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
+                loss, loss_items = compute_loss(pred, targets.to('cpu'))  # loss scaled by batch_size
                 if rank != -1:
                     loss *= opt.world_size  # gradient averaged between devices in DDP mode
                 if opt.quad:
diff --git a/utils/datasets.py b/utils/datasets.py
index ec597b6..de07ff1 100644
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -408,7 +408,7 @@ class LoadImagesAndLabels(Dataset):  # for training/testing
                 x[:, 0] = 0
 
         n = len(shapes)  # number of images
-        bi = np.floor(np.arange(n) / batch_size).astype(np.int)  # batch index
+        bi = np.floor(np.arange(n) / batch_size).astype(np.int64)  # batch index
         nb = bi[-1] + 1  # number of batches
         self.batch = bi  # batch index of image
         self.n = n
@@ -436,7 +436,7 @@ class LoadImagesAndLabels(Dataset):  # for training/testing
                 elif mini > 1:
                     shapes[i] = [1, 1 / mini]
 
-            self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
+            self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int64) * stride
 
         # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
         self.imgs = [None] * n
@@ -648,7 +648,7 @@ def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
     hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
     dtype = img.dtype  # uint8
 
-    x = np.arange(0, 256, dtype=np.int16)
+    x = np.arange(0, 256, dtype=np.int64)
     lut_hue = ((x * r[0]) % 180).astype(dtype)
     lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
     lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
@@ -1034,7 +1034,7 @@ def extract_boxes(path='../coco128/'):  # from utils.datasets import *; extract_
                     b = x[1:] * [w, h, w, h]  # box
                     # b[2:] = b[2:].max()  # rectangle to square
                     b[2:] = b[2:] * 1.2 + 3  # pad
-                    b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int)
+                    b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int64)
 
                     b[[0, 2]] = np.clip(b[[0, 2]], 0, w)  # clip boxes outside of image
                     b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
diff --git a/utils/general.py b/utils/general.py
index 5482629..4cc691b 100644
--- a/utils/general.py
+++ b/utils/general.py
@@ -219,7 +219,7 @@ def labels_to_class_weights(labels, nc=80):
         return torch.Tensor()
 
     labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
-    classes = labels[:, 0].astype(np.int)  # labels = [class xywh]
+    classes = labels[:, 0].astype(np.int64)  # labels = [class xywh]
     weights = np.bincount(classes, minlength=nc)  # occurrences per class
 
     # Prepend gridpoint count (for uCE training)
@@ -234,7 +234,7 @@ def labels_to_class_weights(labels, nc=80):
 
 def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
     # Produces image weights based on class_weights and image contents
-    class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
+    class_counts = np.array([np.bincount(x[:, 0].astype(np.int64), minlength=nc) for x in labels])
     image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
     # index = random.choices(range(n), weights=image_weights, k=1)  # weight image sample
     return image_weights
diff --git a/utils/loss.py b/utils/loss.py
index 9e78df1..2c6d94b 100644
--- a/utils/loss.py
+++ b/utils/loss.py
@@ -164,7 +164,7 @@ class ComputeLoss:
         # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
         na, nt = self.na, targets.shape[0]  # number of anchors, targets
         tcls, tbox, indices, anch = [], [], [], []
-        gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain
+        gain = torch.ones(7, device=targets.device).long()  # normalized to gridspace gain
         ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
         targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices
 

4 训练

cpu

python ./train.py --data ./custom.yaml --cfg ./v5Lite-s-custom.yaml --weights ./v5lite-s.pt --batch-size 64 --device cpu --epochs 100

gpu

python ./train.py --data ./custom.yaml --cfg ./v5Lite-s-custom.yaml --weights ./v5lite-s.pt --batch-size 64 --device 0 --epochs 100

以我的例子，效果如下：

• 2K训练数据，能做到mAP@.5 = 0.67 /mAP.05:.95 =  0.384，大概45轮就没法提升了
• 8K + 10%负向训练数据，能在3个Epoch就做到 0.574 /  0.274，最终0.695 / 0.411，一个Epoch 47s （L4）
• 8K + 50%负向训练数据，

5 模型转换

TODO