pytorch YOLOv4
1.0.0
Uma implementação mínima de pytorch do Yolov4.
Papel Yolo v4: https://arxiv.org/abs/2004.10934
Código fonte: https: //github.com/alexeyab/darknet
Mais detalhes: http://pjreddie.com/darknet/yolo/
Inferência
Trem
├── README.md
├── dataset.py dataset
├── demo.py demo to run pytorch --> tool/darknet2pytorch
├── demo_darknet2onnx.py tool to convert into onnx --> tool/darknet2pytorch
├── demo_pytorch2onnx.py tool to convert into onnx
├── models.py model for pytorch
├── train.py train models.py
├── cfg.py cfg.py for train
├── cfg cfg --> darknet2pytorch
├── data
├── weight --> darknet2pytorch
├── tool
│ ├── camera.py a demo camera
│ ├── coco_annotation.py coco dataset generator
│ ├── config.py
│ ├── darknet2pytorch.py
│ ├── region_loss.py
│ ├── utils.py
│ └── yolo_layer.py
Você pode usar o DarkNet2Pytorch para convertê -lo ou baixar meu modelo convertido.
Use Yolov4 para treinar seus próprios dados
Baixar peso
Transformar dados
Para o conjunto de dados Coco, você pode usar a ferramenta/coco_annotation.py.
# train.txt
image_path1 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
image_path2 x1,y1,x2,y2,id x1,y1,x2,y2,id x1,y1,x2,y2,id ...
...
...
Trem
Você pode definir parâmetros em cfg.py.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
Os modelos Onnx e Tensorrt são convertidos de Pytorch (Tianxiaomo): Pytorch-> Onnx-> Tensorrt. Consulte as seções a seguir para obter mais detalhes das conversões.
| Tipo de modelo | AP | AP50 | AP75 | APS | APM | Apl |
|---|---|---|---|---|---|---|
| Darknet (papel Yolov4) | 0,471 | 0,710 | 0,510 | 0,278 | 0,525 | 0,636 |
| Pytorch (Tianxiaomo) | 0,466 | 0,704 | 0,505 | 0,267 | 0,524 | 0,629 |
| Tensorrt fp32 + lotenmsplugin | 0,472 | 0,708 | 0,511 | 0,273 | 0,530 | 0,637 |
| Tensorrt fp16 + lotenmsplugin | 0,472 | 0,708 | 0,511 | 0,273 | 0,530 | 0,636 |
| Tipo de modelo | AP | AP50 | AP75 | APS | APM | Apl |
|---|---|---|---|---|---|---|
| Darknet (papel Yolov4) | 0,412 | 0,628 | 0,443 | 0,204 | 0,444 | 0,560 |
| Pytorch (Tianxiaomo) | 0,404 | 0,615 | 0,436 | 0,196 | 0,438 | 0,552 |
| Tensorrt fp32 + lotenmsplugin | 0,412 | 0,625 | 0,445 | 0,200 | 0,446 | 0,564 |
| Tensorrt fp16 + lotenmsplugin | 0,412 | 0,625 | 0,445 | 0,200 | 0,446 | 0,563 |
O tamanho da entrada da imagem não é restrito em 320 * 320 , 416 * 416 , 512 * 512 e 608 * 608 . Você pode ajustar seus tamanhos de entrada para uma taxa de entrada diferente, por exemplo: 320 * 608 . O tamanho de entrada maior pode ajudar a detectar alvos menores, mas pode ser mais lento e a memória da GPU exaustão.
height = 320 + 96 * n , n in { 0 , 1 , 2 , 3 , ...}
width = 320 + 96 * m , m in { 0 , 1 , 2 , 3 , ...}Carregue o modelo DarkNet pré -treinado e os pesos da DarkNet para fazer a inferência (o tamanho da imagem já está configurado no arquivo CFG)
python demo.py -cfgfile < cfgFile > -weightfile < weightFile > -imgfile < imgFile >Carregue pesos pytorch (arquivo PTH) para fazer a inferência
python models.py < num_classes > < weightfile > < imgfile > < IN_IMAGE_H > < IN_IMAGE_W > < namefile(optional) >Carregar o arquivo onnx convertido para fazer inferência (consulte a seção 3 e 4)
Carregar o arquivo do motor Tensorrt convertido para fazer inferência (consulte a Seção 5)
Existem 2 saídas de inferência.
[batch, num_boxes, 1, 4] , que representa x1, y1, x2, y2 de cada caixa delimitadora.[batch, num_boxes, num_classes] indicando pontuações de todas as classes para cada caixa delimitadora.Até agora, ainda é necessário um pequeno pedaço de pós-processamento, incluindo o NMS. Estamos tentando minimizar o tempo e a complexidade do pós-processamento.
Este script é para converter o modelo Darknet pré -traido oficial em ONNX
Versão Pytorch Recomendada:
Instale o OnnxRuntime
pip install onnxruntimeExecute o script python para gerar modelo ONNX e executar a demonstração
python demo_darknet2onnx.py < cfgFile > < namesFile > < weightFile > < imageFile > < batchSize > Você pode converter seu modelo Pytorch treinado em ONNX usando este script
Versão Pytorch Recomendada:
Instale o OnnxRuntime
pip install onnxruntimeExecute o script python para gerar modelo ONNX e executar a demonstração
python demo_pytorch2onnx.py < weight_file > < image_path > < batch_size > < n_classes > < IN_IMAGE_H > < IN_IMAGE_W >Por exemplo:
python demo_pytorch2onnx.py yolov4.pth dog.jpg 8 80 416 416Execute o seguinte comando para converter o modelo YOLOV4 ONNX em motor Tensorrt
trtexec --onnx= < onnx_file > --explicitBatch --saveEngine= < tensorRT_engine_file > --workspace= < size_in_megabytes > --fp16Execute o seguinte comando para converter o modelo YOLOV4 ONNX em motor Tensorrt
trtexec --onnx= < onnx_file >
--minShapes=input: < shape_of_min_batch > --optShapes=input: < shape_of_opt_batch > --maxShapes=input: < shape_of_max_batch >
--workspace= < size_in_megabytes > --saveEngine= < engine_file > --fp16Por exemplo:
trtexec --onnx=yolov4_-1_3_320_512_dynamic.onnx
--minShapes=input:1x3x320x512 --optShapes=input:4x3x320x512 --maxShapes=input:8x3x320x512
--workspace=2048 --saveEngine=yolov4_-1_3_320_512_dynamic.engine --fp16python demo_trt.py < tensorRT_engine_file > < input_image > < input_H > < input_W >Essa demonstração aqui funciona apenas quando o BatchSize é dinâmico (1 deve estar dentro da faixa dinâmica) ou em lote = 1, mas você pode atualizar uma demonstração um pouco para outros tamanhos de lote dinâmico ou estático.
Nota1: input_h e input_w devem concordar com o tamanho da entrada no arquivo ONNX original.
Nota2: Operações extras do NMS são necessárias para a saída Tensorrt. Esta demonstração usa o código Python NMS do tool/utils.py .
Primeiro: conversão para ONNX
tensorflow> = 2.0
1: Obrigado: github: https: //github.com/onnx/onnx-tensorflow
2: Execute clone git https://github.com/onnx/onnx-tensorflow.git && cd onnx-tensorflow run pip install -e.
NOTA: Os erros ocorrerão ao usar o "PIP Install Onnx-TF", pelo menos para mim, é recomendável usar a instalação do código-fonte
cd DeepStream
make
Para um único lote,
trtexec --onnx=<onnx_file> --explicitBatch --saveEngine=<tensorRT_engine_file> --workspace=<size_in_megabytes> --fp16
Para vários lojas,
trtexec --onnx=<onnx_file> --explicitBatch --shapes=input:Xx3xHxW --optShapes=input:Xx3xHxW --maxShapes=input:Xx3xHxW --minShape=input:1x3xHxW --saveEngine=<tensorRT_engine_file> --fp16
Nota: as maxshapes não poderiam ser maiores que a forma original do modelo.
Referência:
@article{yolov4,
title={YOLOv4: YOLOv4: Optimal Speed and Accuracy of Object Detection},
author={Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao},
journal = {arXiv},
year={2020}
}
