From 95713b8594e7cda8df502c7762a625764098006d Mon Sep 17 00:00:00 2001
From: wangguanzhong <jerrywgz@126.com>
Date: Mon, 18 May 2020 12:00:00 +0800
Subject: [PATCH] fix ext_op (#711)

---
 ppdet/ext_op/README.md                     |  7 ++-----
 ppdet/ext_op/test/test_corner_pool.py      |  9 ++++++++-
 ppdet/modeling/anchor_heads/corner_head.py | 15 ++++++++++-----
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/ppdet/ext_op/README.md b/ppdet/ext_op/README.md
index 6bfeb0e88..0516cc702 100644
--- a/ppdet/ext_op/README.md
+++ b/ppdet/ext_op/README.md
@@ -59,12 +59,9 @@ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:`python -c 'import paddle; print(paddle.
 执行下列单测，确保自定义算子可在网络中正确使用：
 
 ```
-# 回到 ext_op 目录，添加 PYTHONPATH
+# 回到 ext_op 目录，运行单测
 cd ..
-export PYTHONPATH=$PYTHONPATH:`pwd`
-
-# 运行单测
-python test/test_corner_op.py
+python test/test_corner_pool.py
 ```
 
 单测运行成功会输出提示信息，如下所示：
diff --git a/ppdet/ext_op/test/test_corner_pool.py b/ppdet/ext_op/test/test_corner_pool.py
index 7d89a25de..211aaaded 100755
--- a/ppdet/ext_op/test/test_corner_pool.py
+++ b/ppdet/ext_op/test/test_corner_pool.py
@@ -17,7 +17,14 @@ from __future__ import print_function
 import unittest
 import numpy as np
 import paddle.fluid as fluid
-import cornerpool_lib
+import os
+import sys
+# add python path of PadleDetection to sys.path
+parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
+if parent_path not in sys.path:
+    sys.path.append(parent_path)
+
+from ppdet.ext_op import cornerpool_lib
 
 
 def bottom_pool_np(x):
diff --git a/ppdet/modeling/anchor_heads/corner_head.py b/ppdet/modeling/anchor_heads/corner_head.py
index 01e1195b0..33a3a55ec 100644
--- a/ppdet/modeling/anchor_heads/corner_head.py
+++ b/ppdet/modeling/anchor_heads/corner_head.py
@@ -243,11 +243,6 @@ class CornerHead(object):
                  ae_threshold=1,
                  num_dets=1000,
                  top_k=100):
-        try:
-            import cornerpool_lib
-        except:
-            logger.error(
-                "cornerpool_lib not found, compile in ppdet/ext_op at first")
         self.train_batch_size = train_batch_size
         self.test_batch_size = test_batch_size
         self.num_classes = num_classes
@@ -279,6 +274,11 @@ class CornerHead(object):
         return conv1
 
     def get_output(self, input):
+        try:
+            from ppdet.ext_op import cornerpool_lib
+        except:
+            logger.error(
+                "cornerpool_lib not found, compile in ppdet/ext_op at first")
         for ind in range(self.stack):
             cnv = input[ind]
             tl_modules = corner_pool(
@@ -455,6 +455,11 @@ class CornerHead(object):
         return {'loss': loss}
 
     def get_prediction(self, input):
+        try:
+            from ppdet.ext_op import cornerpool_lib
+        except:
+            logger.error(
+                "cornerpool_lib not found, compile in ppdet/ext_op at first")
         ind = self.stack - 1
         tl_modules = corner_pool(
             input,
-- 
GitLab