Nota especial: consulte el código oficial yolov8 de código abierto, los documentos oficiales de Rockchip y los documentos oficiales de Horizon. Si hay alguna infracción, infórmela y elimínela, gracias.
El modelo, la imagen de prueba, el resultado de la prueba y el código completo se colocan en github, y el modelo y el código del enlace de referencia .
Desde que escribí varias publicaciones de blog relacionadas con la implementación de chips en el lado de la placa relacionadas con la detección y segmentación de yolov8, algunos internautas me pidieron que escribiera un blog sobre la implementación de yolov8pose, y la implementación de yolov8pose está aquí.
Nota especial : en este ejemplo, el entrenamiento del modelo no utiliza muchos datos y no se puede garantizar el efecto del modelo. Solo se usa para pruebas e implementación. Si cambia una imagen, es posible que no se detecte, lo cual es un Fenómeno normal.
1 Modelo y formación
El código de entrenamiento se refiere al código de entrenamiento oficial de código abierto yolov8seg. Dado que SiLU no es compatible con algunos chips del lado de la placa, se cambia a ReLU. El conjunto de datos de entrenamiento es parte del entrenamiento coco. Se utiliza principalmente para probar el proceso. y no se puede lograr el efecto del modelo. Garantizado, otra prueba de imagen no detectará la normalidad.
2 Exportar onnx de yolov8pose
El método de exportación de onnx proporcionado en este ejemplo solo es adecuado para el código del almacén correspondiente de este ejemplo. Si utiliza el onnx exportado oficialmente, escriba el código de posprocesamiento usted mismo. Gracias
Algunos cálculos en el posprocesamiento son ineficientes o no son compatibles con el chip del lado de la placa. Para exportar onnx, debe evitar los operadores que no son amigables o no son compatibles con el chip del lado de la placa. Exporte las partes modificadas de onnx.
Paso 1:
Haga predicciones y guarde solo los pesos de pt. Agregue el código como se muestra a continuación.
# 保存权重值
import torch
self.model.fuse()
self.model.eval()
torch.save(self.model.state_dict(), './weights/yolov8pos_relu_dict.pt')
Ejecute el siguiente código después de la modificación (yolov8pos_relu_dict.pt se genera en la carpeta de pesos):
# 推理
model = YOLO('./weights/yolov8pos_relu.pt')
results = model(task='pose', mode='predict', source='./images/test.jpg', line_width=3, show=True, save=True, device='cpu')
Paso 2:
exporte onnx y elimine los operadores innecesarios. Modifique el código de la siguiente manera.
(1) Modificar el cabezal de detección
Modificar el código completo del módulo:
class Detect(nn.Module):
"""YOLOv8 Detect head for detection models."""
dynamic = False # force grid reconstruction
export = False # export mode
shape = None
anchors = torch.empty(0) # init
strides = torch.empty(0) # init
def __init__(self, nc=80, ch=()): # detection layer
super().__init__()
self.nc = nc # number of classes
self.nl = len(ch) # number of detection layers
self.reg_max = 16 # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
self.no = nc + self.reg_max * 4 # number of outputs per anchor
self.stride = torch.zeros(self.nl) # strides computed during build
c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], self.nc) # channels
self.cv2 = nn.ModuleList(
nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
# 导出 onnx 增加
self.conv1x1 = nn.Conv2d(16, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(16, dtype=torch.float)
self.conv1x1.weight.data[:] = nn.Parameter(x.view(1, 16, 1, 1))
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
shape = x[0].shape # BCHW
y = []
for i in range(self.nl):
t1 = self.cv2[i](x[i])
t2 = self.cv3[i](x[i])
y.append(self.conv1x1(t1.view(t1.shape[0], 4, 16, -1).transpose(2, 1).softmax(1)))
# y.append(t2.sigmoid())
y.append(t2)
return y
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.training:
return x
elif self.dynamic or self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'): # avoid TF FlexSplitV ops
box = x_cat[:, :self.reg_max * 4]
cls = x_cat[:, self.reg_max * 4:]
else:
box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
y = torch.cat((dbox, cls.sigmoid()), 1) # 官方代码
# y = torch.cat((self.dfl(box), cls.sigmoid()), 1) # 导出本实例的onnx使用
return y if self.export else (y, x)
def bias_init(self):
"""Initialize Detect() biases, WARNING: requires stride availability."""
m = self # self.model[-1] # Detect() module
# cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
# ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum()) # nominal class frequency
for a, b, s in zip(m.cv2, m.cv3, m.stride): # from
a[-1].bias.data[:] = 1.0 # box
b[-1].bias.data[:m.nc] = math.log(5 / m.nc / (640 / s) ** 2) # cls (.01 objects, 80 classes, 640 img)
(2) Modificar el encabezado de pose
class Pose(Detect):
"""YOLOv8 Pose head for keypoints models."""
def __init__(self, nc=80, kpt_shape=(17, 3), ch=()):
"""Initialize YOLO network with default parameters and Convolutional Layers."""
super().__init__(nc, ch)
self.kpt_shape = kpt_shape # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
self.nk = kpt_shape[0] * kpt_shape[1] # number of keypoints total
self.detect = Detect.forward
c4 = max(ch[0] // 4, self.nk)
self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nk, 1)) for x in ch)
def forward(self, x):
"""Perform forward pass through YOLO model and return predictions."""
bs = x[0].shape[0] # batch size
# kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1) # (bs, 17*3, h*w)
ps = []
for i in range(self.nl):
ps.append(self.cv4[i](x[i]))
x = self.detect(self, x)
return x, ps
if self.training:
return x, kpt
pred_kpt = self.kpts_decode(bs, kpt)
return torch.cat([x, pred_kpt], 1) if self.export else (torch.cat([x[0], pred_kpt], 1), (x[1], kpt))
def kpts_decode(self, bs, kpts):
"""Decodes keypoints."""
ndim = self.kpt_shape[1]
if self.export: # required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
y = kpts.view(bs, *self.kpt_shape, -1)
a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
if ndim == 3:
a = torch.cat((a, y[:, :, 1:2].sigmoid()), 2)
return a.view(bs, self.nk, -1)
else:
y = kpts.clone()
if ndim == 3:
y[:, 2::3].sigmoid_() # inplace sigmoid
y[:, 0::ndim] = (y[:, 0::ndim] * 2.0 + (self.anchors[0] - 0.5)) * self.strides
y[:, 1::ndim] = (y[:, 1::ndim] * 2.0 + (self.anchors[1] - 0.5)) * self.strides
return y
(3) Agregue el código para generar onnx
def _new(self, cfg: str, task=None, verbose=True):
"""
Initializes a new model and infers the task type from the model definitions.
Args:
cfg (str): model configuration file
task (str) or (None): model task
verbose (bool): display model info on load
"""
cfg_dict = yaml_model_load(cfg)
self.cfg = cfg
self.task = task or guess_model_task(cfg_dict)
self.model = TASK_MAP[self.task][0](cfg_dict, verbose=verbose and RANK == -1) # build model
self.overrides['model'] = self.cfg
# Below added to allow export from yamls
args = {
**DEFAULT_CFG_DICT, **self.overrides} # combine model and default args, preferring model args
self.model.args = {
k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS} # attach args to model
self.model.task = self.task
import torch
self.model.fuse()
self.model.eval()
self.model.load_state_dict(torch.load('./weights/yolov8pos_relu_dict.pt', map_location='cpu'), strict=False)
print("=========== onnx =========== ")
dummy_input = torch.randn(1, 3, 640, 640)
input_names = ["data"]
output_names = ["cls1", "reg1", "cls2", "reg2", "cls3", "reg3", "ps1", "ps2", "ps3"]
torch.onnx.export(self.model, dummy_input, "./weights/yolov8pos_relu.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=11)
print("======================== convert onnx Finished! .... ")
Después de modificar los tres bloques anteriores, ejecute el siguiente código (para generar el modelo onnx):
model = YOLO('./ultralytics/models/v8/yolov8-pose.yaml')
Nota: La siguiente secuencia de modificación no debe ser incorrecta, los códigos de las dos ejecuciones son diferentes, preste atención, preste atención, preste atención.
3 resultados de las pruebas de pytoch y onnx de yolov8pose
(1) resultados de la prueba de pytorch
(2) resultados de la prueba de onnx
4 Referencia de prueba de simulación de horizonte y horizonte Rockchip rknn
La construcción del entorno Rockchip y los pasos detallados se refieren a [Conversión del modelo Rockchip RKNN y simulación de PC] .
Para obtener detalles sobre cómo configurar el entorno de Horizon y los pasos detallados, consulte [Conversión del modelo de Horizon y prueba de simulación en la PC] .
5 enlaces relacionados
yolov8 detecta el despliegue de prueba de simulación de chips Rockchip RKNN y Horizon Horizon