Träna en objektidentifieringsmodell med TensorFlow

Nu när vi har konfigurerat TensorFlow använder vi YOLO-arkitekturen för att träna objektidentifieringsmodellen. YOLO är ett neuralt nätverk som förutsäger avgränsningsrutor och klassannolikheter från en bild i en enda utvärdering. YOLO-modeller kan bearbeta över 60 bildrutor per sekund, vilket gör det till en bra arkitektur för att identifiera objekt i videor. Du hittar mer information om hur YOLO fungerar här.

Använda YOLO

Börja med att ladda ned den här YOLO-exempelfilen, som innehåller hjälpskript för att komma igång.

När du använder YOLO har vi tre alternativ:

1. Använd förtränade modellvikter för YOLO. Den förtränade modellen har tränats på en stor datamängd med 80 klasser (kategorier) för vardagliga föremål som buss, person, smörgås osv. Om du vill ladda ned en förtränad YOLO-modell i ONNX-format kan du göra det här. Sedan kan du gå vidare till det sista steget i den här självstudien för att lära dig hur du integrerar modellen i en app.

2. Implementera överföringsinlärning med en anpassad datauppsättning. Överföringsinlärning är en metod för att använda en tränad modell som utgångspunkt för att träna en modell som löser en annan men relaterad uppgift. Den här handledningen använder förtränade vikter för YOLO med 80 klasser för att träna en modell med 20 klasser med hjälp av VOC-datamängden. Om du vill skapa en egen datauppsättning med anpassade klasser kan du läsa anvisningarna här.

3. Träna YOLO från grunden. Den här tekniken rekommenderas inte eftersom det är mycket svårt att konvergera. Det ursprungliga YOLO-papperet tränade darknet på ImageNet (som innehåller hundratusentals foton) innan hela nätverket också tränades.

Implementera överföringsinlärning på förtränade YOLO-vikter till VOC-datamängden:

Vi fortsätter med det andra alternativet och implementerar överföringsinlärning med följande steg.

1. I ett miniconda-fönster navigerar du till yolo-exempelkatalogen och kör följande kommando för att installera alla nödvändiga pip-paket för YOLO.

pip install -r requirements.txt

2. Kör installationsskriptet för att ladda ned data och förtränade vikter

python setup.py

3. Transformera datauppsättningen. Se tools/voc2012.py för implementering – det här formatet baseras på tensorflow-objektidentifierings-API:et. Många fält krävs inte, men här har fyllts i för kompatibilitet med det officiella API:et.

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split train \
  --output_file ./data/voc2012_train.tfrecord

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split val \
  --output_file ./data/voc2012_val.tfrecord

4. Träna modellen. Kör följande kommandon:

python convert.py
python detect.py --image ./data/meme.jpg # Sanity check

python train.py \
	--dataset ./data/voc2012_train.tfrecord \
	--val_dataset ./data/voc2012_val.tfrecord \
	--classes ./data/voc2012.names \
	--num_classes 20 \
	--mode fit --transfer darknet \
	--batch_size 16 \
	--epochs 10 \
	--weights ./checkpoints/yolov3.tf \
	--weights_num_classes 80 

Nu har du en omtränad modell med 20 klasser som är redo att användas.

Nästa steg

Nu när vi har skapat en TensorFlow-modell måste vi konvertera den till ONNX-format för användning med Windows Machine Learning-API:erna.