未验证 提交 79b64608 编写于 作者: W Wang Feng 提交者: GitHub

chore(gpu_nms): use a better implemention (#51)

上级 dd76c1f9
...@@ -11,24 +11,23 @@ import megengine.functional as F ...@@ -11,24 +11,23 @@ import megengine.functional as F
from megengine._internal.craniotome import CraniotomeBase from megengine._internal.craniotome import CraniotomeBase
from megengine.core.tensor import wrap_io_tensor from megengine.core.tensor import wrap_io_tensor
_so_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lib_nms.so") _current_path = os.path.dirname(os.path.abspath(__file__))
_so_path = os.path.join(_current_path, "lib_nms.so")
try: try:
_so_lib = ctypes.CDLL(_so_path) _so_lib = ctypes.CDLL(_so_path)
except Exception: except Exception:
MGE = mge.__file__.rsplit("/", 1)[0] import subprocess
current_path = os.path.abspath(__file__).rsplit("/", 1)[0] mge_path = os.path.join(os.path.dirname(mge.__file__), "_internal", "include")
mge_path = os.path.join(MGE, "_internal/include") assert os.path.exists(mge_path), "{} file not found".format(mge_path)
src_file = os.path.join(current_path, "gpu_nms/nms.cu") src_file = os.path.join(_current_path, "gpu_nms", "nms.cu")
dst_file = os.path.join(current_path, "lib_nms.so") assert os.path.exists(src_file), "{} file not found".format(src_file)
assert os.path.exists(mge_path)
assert os.path.exists(src_file)
cmd = ( cmd = (
"nvcc -I {} -shared -o {} -Xcompiler '-fno-strict-aliasing -fPIC' {}".format( "nvcc -I {} -shared -o {} -Xcompiler '-fno-strict-aliasing -fPIC' {}".format(
mge_path, dst_file, src_file mge_path, _so_path, src_file
) )
) )
os.system(cmd) subprocess.check_call(cmd, shell=True)
_so_lib = ctypes.CDLL(_so_path) _so_lib = ctypes.CDLL(_so_path)
_TYPE_POINTER = ctypes.c_void_p _TYPE_POINTER = ctypes.c_void_p
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册