Robert001's picture
first commit
b334e29
raw
history blame
2.02 kB
# Copyright (c) OpenMMLab. All rights reserved.
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