Import and convert FP16 weights from TensorFlow

modules/dnn/src/tensorflow/tf_importer.cpp

@@ -63,10 +63,15 @@ void blobShapeFromTensor(const tensorflow::TensorProto &tensor, MatShape& shape)
     {
         const tensorflow::TensorShapeProto &_shape = tensor.tensor_shape();
         int i, n = _shape.dim_size();
-        shape.resize(n);
+        if (n)
+        {
+            shape.resize(n);
 
-        for (i = 0; i < n; i++)
-            shape[i] = (int)_shape.dim(i).size();
+            for (i = 0; i < n; i++)
+                shape[i] = (int)_shape.dim(i).size();
+        }
+        else
+            shape.resize(1, 1);  // Scalar.
     }
     else
     {
@@ -74,6 +79,43 @@ void blobShapeFromTensor(const tensorflow::TensorProto &tensor, MatShape& shape)
     }
 }
 
+static Mat getTensorContent(const tensorflow::TensorProto &tensor)
+{
+    std::string content = tensor.tensor_content();
+    switch (tensor.dtype())
+    {
+        case tensorflow::DT_FLOAT:
+            return Mat(1, content.size() / sizeof(float), CV_32FC1, (void*)content.c_str()).clone();
+        case tensorflow::DT_DOUBLE:
+            return Mat(1, content.size() / sizeof(double), CV_64FC1, (void*)content.c_str()).clone();
+        case tensorflow::DT_HALF:
+        {
+            Mat halfs;
+            if (!content.empty())
+            {
+                static const int kHalfSize = 2;
+                halfs = Mat(1, content.size() / kHalfSize, CV_16UC1, (void*)content.c_str());
+            }
+            else
+            {
+                const RepeatedField<int32_t>& field = tensor.half_val();
+                CV_Assert(!field.empty());
+                Mat ints(1, field.size(), CV_32SC1, (void*)field.data());
+                ints.convertTo(halfs, CV_16UC1);
+            }
+            // Reinterpret as a signed shorts just for a convertFp16 call.
+            Mat halfsSigned(halfs.size(), CV_16SC1, halfs.data);
+            Mat floats(halfs.size(), CV_32FC1);
+            convertFp16(halfsSigned, floats);
+            return floats;
+        }
+        default:
+            CV_Error(Error::StsError, "Tensor's data type is not supported");
+            break;
+    }
+    return Mat();
+}
+
 template <typename T>
 void parseTensor(const tensorflow::TensorProto &tensor, Mat &dstBlob)
 {
@@ -90,11 +132,12 @@ void parseTensor(const tensorflow::TensorProto &tensor, Mat &dstBlob)
 
     dstBlob.create(shape, CV_32F);
 
-    int size = tensor.tensor_content().size() / sizeof(T);
+    Mat tensorContent = getTensorContent(tensor);
+    int size = tensorContent.total();
     CV_Assert(size == (int)dstBlob.total());
 
     float *dstData = dstBlob.ptr<float>();
-    const T *data = reinterpret_cast<const T*>(tensor.tensor_content().c_str());
+    const T *data = reinterpret_cast<const T*>(tensorContent.data);
 
     if (dims == 4)
     {
@@ -125,6 +168,7 @@ void blobFromTensor(const tensorflow::TensorProto &tensor, Mat &dstBlob)
 {
     switch (tensor.dtype()) {
         case tensorflow::DT_FLOAT:
+        case tensorflow::DT_HALF:
             parseTensor<float>(tensor, dstBlob);
             break;
         case tensorflow::DT_DOUBLE:
@@ -406,7 +450,8 @@ void TFImporter::kernelFromTensor(const tensorflow::TensorProto &tensor, Mat &ds
     int dims = (int)shape.size();
 
     // TODO: other blob types
-    CV_Assert(tensor.dtype() == tensorflow::DT_FLOAT);
+    CV_Assert(tensor.dtype() == tensorflow::DT_FLOAT ||
+              tensor.dtype() == tensorflow::DT_HALF);
     CV_Assert(dims == 4);
 
     // REORDER kernel HWIO to OIHW
@@ -416,11 +461,12 @@ void TFImporter::kernelFromTensor(const tensorflow::TensorProto &tensor, Mat &ds
 
     dstBlob.create(shape, CV_32F);
 
-    int size = tensor.tensor_content().size() / sizeof(float);
+    Mat tensorContent = getTensorContent(tensor);
+    int size = tensorContent.total();
     CV_Assert(size == (int)dstBlob.total());
 
     float *dstData = dstBlob.ptr<float>();
-    const float *data = reinterpret_cast<const float*>(tensor.tensor_content().c_str());
+    const float *data = reinterpret_cast<const float*>(tensorContent.data);
 
     int out_c = shape[0], input_c = shape[1], height = shape[2], width = shape[3];
     int total = out_c*input_c*height*width;
@@ -753,7 +799,16 @@ void TFImporter::populateNet(Net dstNet)
                 // Multiplication by constant.
                 CV_Assert(layer.input_size() == 2);
 
-                float scale = getConstBlob(layer, value_id).float_val()[0];
+                float scale;
+                if (!getConstBlob(layer, value_id).float_val().empty())
+                    scale = getConstBlob(layer, value_id).float_val()[0];
+                else
+                {
+                    Mat scaleMat;
+                    blobFromTensor(getConstBlob(layer, value_id), scaleMat);
+                    CV_Assert(scaleMat.total() == 1 && scaleMat.type() == CV_32FC1);
+                    scale = scaleMat.at<float>(0, 0);
+                }
                 layerParams.set("scale", scale);
 
                 int id = dstNet.addLayer(name, "Power", layerParams);

modules/dnn/test/test_tf_importer.cpp

@@ -76,7 +76,8 @@ static std::string path(const std::string& file)
     return findDataFile("dnn/tensorflow/" + file, false);
 }
 
-static void runTensorFlowNet(const std::string& prefix)
+static void runTensorFlowNet(const std::string& prefix,
+                             double l1 = 1e-5, double lInf = 1e-4)
 {
     std::string netPath = path(prefix + "_net.pb");
     std::string inpPath = path(prefix + "_in.npy");
@@ -89,7 +90,7 @@ static void runTensorFlowNet(const std::string& prefix)
 
     net.setInput(input);
     cv::Mat output = net.forward();
-    normAssert(target, output);
+    normAssert(target, output, "", l1, lInf);
 }
 
 TEST(Test_TensorFlow, single_conv)
@@ -130,4 +131,19 @@ TEST(Test_TensorFlow, deconvolution)
     runTensorFlowNet("deconvolution");
 }
 
+TEST(Test_TensorFlow, fp16)
+{
+    const float l1 = 1e-3;
+    const float lInf = 1e-2;
+    runTensorFlowNet("fp16_single_conv", l1, lInf);
+    runTensorFlowNet("fp16_deconvolution", l1, lInf);
+    runTensorFlowNet("fp16_max_pool_odd_same", l1, lInf);
+    runTensorFlowNet("fp16_padding_valid", l1, lInf);
+    runTensorFlowNet("fp16_eltwise_add_mul", l1, lInf);
+    runTensorFlowNet("fp16_max_pool_odd_valid", l1, lInf);
+    runTensorFlowNet("fp16_pad_and_concat", l1, lInf);
+    runTensorFlowNet("fp16_max_pool_even", l1, lInf);
+    runTensorFlowNet("fp16_padding_same", l1, lInf);
+}
+
 }