convert_from_tensorflow.py: add support when kernel size is 1*1 with one input/output channel (gray image)

Guo, Yejun committed 6 years ago

Commit e52070e1 parent ed9fc2e1 file · +4 −1Browse files

Showing the changes to tools/python/convert_from_tensorflow.py (1 of 1 file). Show all

Unified · Split

         filter_width = ktensor.tensor_shape.dim[1].size
         in_channels = ktensor.tensor_shape.dim[2].size
         out_channels = ktensor.tensor_shape.dim[3].size
-        kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
+        if filter_height * filter_width * in_channels * out_channels == 1:
+            kernel = np.float32(ktensor.float_val[0])
+        else:
+            kernel = np.frombuffer(ktensor.tensor_content, dtype=np.float32)
         kernel = kernel.reshape(filter_height, filter_width, in_channels, out_channels)
         kernel = np.transpose(kernel, [3, 0, 1, 2])

Keyboard shortcuts