|
|
|
import importlib |
|
import os |
|
import pkgutil |
|
import warnings |
|
from collections import namedtuple |
|
|
|
import torch |
|
|
|
if torch.__version__ != 'parrots': |
|
|
|
def load_ext(name, funcs): |
|
ext = importlib.import_module('mmcv.' + name) |
|
for fun in funcs: |
|
assert hasattr(ext, fun), f'{fun} miss in module {name}' |
|
return ext |
|
else: |
|
from parrots import extension |
|
from parrots.base import ParrotsException |
|
|
|
has_return_value_ops = [ |
|
'nms', |
|
'softnms', |
|
'nms_match', |
|
'nms_rotated', |
|
'top_pool_forward', |
|
'top_pool_backward', |
|
'bottom_pool_forward', |
|
'bottom_pool_backward', |
|
'left_pool_forward', |
|
'left_pool_backward', |
|
'right_pool_forward', |
|
'right_pool_backward', |
|
'fused_bias_leakyrelu', |
|
'upfirdn2d', |
|
'ms_deform_attn_forward', |
|
'pixel_group', |
|
'contour_expand', |
|
] |
|
|
|
def get_fake_func(name, e): |
|
|
|
def fake_func(*args, **kwargs): |
|
warnings.warn(f'{name} is not supported in parrots now') |
|
raise e |
|
|
|
return fake_func |
|
|
|
def load_ext(name, funcs): |
|
ExtModule = namedtuple('ExtModule', funcs) |
|
ext_list = [] |
|
lib_root = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) |
|
for fun in funcs: |
|
try: |
|
ext_fun = extension.load(fun, name, lib_dir=lib_root) |
|
except ParrotsException as e: |
|
if 'No element registered' not in e.message: |
|
warnings.warn(e.message) |
|
ext_fun = get_fake_func(fun, e) |
|
ext_list.append(ext_fun) |
|
else: |
|
if fun in has_return_value_ops: |
|
ext_list.append(ext_fun.op) |
|
else: |
|
ext_list.append(ext_fun.op_) |
|
return ExtModule(*ext_list) |
|
|
|
|
|
def check_ops_exist(): |
|
ext_loader = pkgutil.find_loader('mmcv._ext') |
|
return ext_loader is not None |
|
|
