| |
| 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 |
|
|
