From 67622f76a3bb183490505068f93ce45c5e5d0b8a Mon Sep 17 00:00:00 2001
From: wang-xinyu <wangxinyu_es@163.com>
Date: Thu, 28 Nov 2019 00:28:33 +0800
Subject: [PATCH] yolov3 add fp16 mode

---
 yolov3/README.md  | 2 ++
 yolov3/yolov3.cpp | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/yolov3/README.md b/yolov3/README.md
index dbe5908..265d984 100644
--- a/yolov3/README.md
+++ b/yolov3/README.md
@@ -13,6 +13,8 @@ Following tricks are used in this yolov3,
 - upsample layer is replaced by a deconvolution layer.
 - Batchnorm layer, implemented by scale layer.
 
+For FP16 mode, just need add one line `builder->setFp16Mode(true);`. On my TX1, it's 115ms in fp16, while 145ms in fp32.
+
 ```
 // 1. generate yolov3.wts from [pytorchx/pytorch-yolo-v3](https://github.com/wang-xinyu/pytorchx/tree/master/pytorch-yolo-v3)
 
diff --git a/yolov3/yolov3.cpp b/yolov3/yolov3.cpp
index 63d05d0..3706d5b 100644
--- a/yolov3/yolov3.cpp
+++ b/yolov3/yolov3.cpp
@@ -12,6 +12,8 @@
 #include "yololayer.h"
 #include <opencv2/opencv.hpp>
 
+#define USE_FP16  // comment out this if want to use FP32
+
 // stuff we know about the network and the input/output blobs
 static const int INPUT_H = 320;
 static const int INPUT_W = 320;
@@ -383,6 +385,9 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType
     // Build engine
     builder->setMaxBatchSize(maxBatchSize);
     builder->setMaxWorkspaceSize(1 << 20);
+#ifdef USE_FP16
+    builder->setFp16Mode(true);
+#endif
     ICudaEngine* engine = builder->buildCudaEngine(*network);
     std::cout << "build out" << std::endl;