def register_extra_symbolics(opset=10):
assert opset >= 10
register_op('addcmul', addcmul_symbolic, '', opset)
register_op('view_as', view_as_symbolic, '', opset)
register_op('topk', topk_symbolic, '', opset)
register_op('group_norm', group_norm_symbolic, '', opset)
register_op('nms_core', nms_core_symbolic, 'mmdet_custom', opset)
def register_quantized_ops(domain: str, version: int):
# Register all the non-quantized ops
symbolic_registry.register_version("", version)
# Register all quantized ops
module = importlib.import_module("torch.onnx.symbolic_caffe2")
symbolic_registry._symbolic_versions["caffe2"] = module
quant_version_ops = inspect.getmembers(
symbolic_registry._symbolic_versions["caffe2"])
for op in quant_version_ops:
if inspect.isfunction(
op[1]) and not symbolic_registry.is_registered_op(
op[0], domain, version):
aten_q_ops = [
"relu",
"_empty_affine_quantized",
"dequantize",
"quantize_per_tensor",
"upsample_nearest2d",
"avg_pool2d",
"reshape",
"slice",
"cat",
"max_pool2d",
"sigmoid",
]
if op[0] in aten_q_ops:
symbolic_registry.register_op(op[0], op[1], "", version)
symbolic_registry.register_op(op[0], op[1], domain, version)
def register_quantized_ops(domain, version):
# Register all the non-quantized ops
sym_registry.register_version('', version)
# Register all quantized ops
module = importlib.import_module('torch.onnx.symbolic_caffe2')
sym_registry._symbolic_versions['caffe2'] = module
quant_version_ops = getmembers(sym_registry._symbolic_versions['caffe2'])
for op in quant_version_ops:
if isfunction(op[1]) and not sym_registry.is_registered_op(op[0], domain, version):
aten_q_ops = ['relu', '_empty_affine_quantized', 'dequantize', 'quantize_per_tensor', 'upsample_nearest2d']
if op[0] in aten_q_ops:
sym_registry.register_op(op[0], op[1], '', version)
sym_registry.register_op(op[0], op[1], domain, version)
def register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version):
if not bool(re.match(r"^[a-zA-Z0-9-_]*::[a-zA-Z]+[a-zA-Z0-9-_]*$", symbolic_name)):
raise RuntimeError("Failed to register operator {}. \
The symbolic name must match the format Domain::Name, \
and sould start with a letter and contain only \
alphanumerical characters"
.format(symbolic_name))
ns, op_name = symbolic_name.split('::')
unaccepted_domain_names = ["onnx", "aten", "prim"]
if ns in unaccepted_domain_names:
raise RuntimeError("Failed to register operator {}. The domain {} is already a used domain."
.format(symbolic_name, ns))
import torch.onnx.symbolic_registry as sym_registry
sym_registry.register_op(op_name, symbolic_fn, ns, opset_version)
def main():
opt = parser.parse_args()
print(torch.__version__)
print(opt)
enc_layers = [1, 2, 2, 4]
dec_layers = [1, 1, 1, 1]
number_of_channels = [
int(8 * 2**i) for i in range(1, 1 + len(enc_layers))
] #[16,32,64,128]
model = UNet(depth=len(enc_layers),
encoder_layers=enc_layers,
decoder_layers=dec_layers,
number_of_channels=number_of_channels,
number_of_outputs=3)
s = torch.load(os.path.join(opt.models_path, opt.name,
opt.name + 'best_model.pth'),
map_location='cpu')
new_state_dict = OrderedDict()
for k, v in s['model'].state_dict().items():
name = k[7:] # remove 'module' word in the beginning of keys.
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.eval()
x = torch.randn(1,
4,
opt.input_size[0],
opt.input_size[1],
opt.input_size[2],
requires_grad=True)
register_op("group_norm", group_norm_symbolic, "", 10)
torch_out = torch.onnx.export(
model, # model being run
[
x,
], # model input (or a tuple for multiple inputs)
os.path.join(
opt.models_path, opt.name, opt.name + ".onnx"
), # where to save the model (can be a file or file-like object)
export_params=True,
verbose=
True, # store the trained parameter weights inside the model file
opset_version=10)
def __init__(self,
inner,
namespace,
name,
verbose=False,
force_rebuild=True,
unique_name=True):
super().__init__()
self.inner = inner
self.namespace = namespace
self.name = name
if unique_name:
self.name = name + '_' + next(
tempfile._get_candidate_names()) + 'x'
self.qualified_name = '{}::{}'.format(self.namespace, self.name)
self.num_outputs = 1
self.params = {}
self.verbose = verbose
self.force_rebuild = force_rebuild
# Register symbolic function for ONNX export.
while not is_registered_op(self.name, self.namespace, 10):
register_op(self.name, self.symbolic, self.namespace, 10)
def export_onnx(model, cfg, output_file_path='model', disable_dyn_axes=True,
verbose=False, opset=9, extra_check=False):
transform = build_inference_transform(
cfg.data.height,
cfg.data.width,
norm_mean=cfg.data.norm_mean,
norm_std=cfg.data.norm_std,
)
input_img = random_image(cfg.data.height, cfg.data.width)
device = next(model.parameters()).device
input_blob = transform(input_img).unsqueeze(0).to(device)
input_names = ['data']
output_names = ['output']
if not disable_dyn_axes:
dynamic_axes = {'data': {0: 'batch_size', 1: 'channels', 2: 'height', 3: 'width'}}
else:
dynamic_axes = {}
if not output_file_path.endswith('.onnx'):
output_file_path += '.onnx'
register_op("group_norm", group_norm_symbolic, "", opset)
register_op("relu6", relu6_symbolic, "", opset)
register_op("hardsigmoid", hardsigmoid_symbolic, "", opset)
with torch.no_grad():
torch.onnx.export(
model,
input_blob,
output_file_path,
verbose=verbose,
export_params=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes,
opset_version=opset,
operator_export_type=torch.onnx.OperatorExportTypes.ONNX
)
if extra_check:
net_from_onnx = onnx.load(output_file_path)
try:
onnx.checker.check_model(net_from_onnx)
print('ONNX check passed.')
except onnx.onnx_cpp2py_export.checker.ValidationError as ex:
print(f'ONNX check failed: {ex}.')
return output_file_path
def register_extra_symbolics(opset=11):
register_op('one_hot', one_hot, '', opset)
register_op('im2col', im2col, '', opset)
register_op('topk', topk, '', opset)
register_op('softmax', softmax, '', opset)
register_op('constant_pad_nd', constant_pad_nd, '', opset)
register_op('reflection_pad1d', reflection_pad1d, '', opset)
register_op('reflection_pad2d', reflection_pad2d, '', opset)
register_op('reflection_pad3d', reflection_pad3d, '', opset)
register_op('avg_pool1d', avg_pool1d, '', opset)
register_op('avg_pool2d', avg_pool2d, '', opset)
register_op('avg_pool3d', avg_pool3d, '', opset)
register_op('adaptive_avg_pool1d', adaptive_avg_pool1d, '', opset)
register_op('adaptive_avg_pool2d', adaptive_avg_pool2d, '', opset)
register_op('adaptive_avg_pool3d', adaptive_avg_pool3d, '', opset)
register_op('masked_select', masked_select, '', opset)
register_op('upsample_nearest1d', upsample_nearest1d, '', opset)
register_op('upsample_nearest2d', upsample_nearest2d, '', opset)
register_op('upsample_nearest3d', upsample_nearest3d, '', opset)
register_op('upsample_linear1d', upsample_linear1d, '', opset)
register_op('upsample_bilinear2d', upsample_bilinear2d, '', opset)
register_op('upsample_trilinear3d', upsample_trilinear3d, '', opset)
register_op('upsample_bicubic2d', upsample_bicubic2d, '', opset)
register_op('new_full', new_full, '', opset)
register_op('grid_sampler', grid_sampler, '', opset)
register_op('cummax', cummax, '', opset)
register_op('cummin', cummin, '', opset)
register_op('roll', roll, '', opset)
def register_extra_symbolics(opset=10):
assert opset >= 10
register_op("addcmul", addcmul_symbolic, "", opset)
register_op("view_as", view_as_symbolic, "", opset)
register_op("topk", topk_symbolic, "", opset)
register_op("group_norm", group_norm_symbolic, "", opset)
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
help='Path to config file')
parser.add_argument('--output-name',
type=str,
default='model',
help='Path to save ONNX model')
parser.add_argument('--opset', type=int, default=9)
parser.add_argument('--verbose',
default=False,
action='store_true',
help='Verbose mode for onnx.export')
parser.add_argument('opts',
default=None,
nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg)
cfg.merge_from_list(args.opts)
cfg.freeze()
num_classes = parse_num_classes(cfg.data.sources)
model = build_model(**model_kwargs(cfg, num_classes))
load_pretrained_weights(model, cfg.model.load_weights)
model.eval()
transform = build_inference_transform(
cfg.data.height,
cfg.data.width,
norm_mean=cfg.data.norm_mean,
norm_std=cfg.data.norm_std,
)
input_img = random_image(cfg.data.height, cfg.data.width)
input_blob = transform(input_img).unsqueeze(0)
input_names = ['data']
output_names = ['reid_embedding']
dynamic_axes = {
'data': {
0: 'batch_size',
1: 'channels',
2: 'height',
3: 'width'
},
'reid_embedding': {
0: 'batch_size',
1: 'dim'
}
}
output_file_path = args.output_name
if not args.output_name.endswith('.onnx'):
output_file_path += '.onnx'
register_op("group_norm", group_norm_symbolic, "", args.opset)
with torch.no_grad():
torch.onnx.export(
model,
input_blob,
output_file_path,
verbose=args.verbose,
export_params=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes,
opset_version=args.opset,
operator_export_type=torch.onnx.OperatorExportTypes.ONNX)
net_from_onnx = onnx.load(output_file_path)
try:
onnx.checker.check_model(net_from_onnx)
print('ONNX check passed.')
except onnx.onnx_cpp2py_export.checker.ValidationError as ex:
print('ONNX check failed: {}.'.format(ex))
def register_extra_symbolics(opset=11):
# Following strings of text style are from colorama package
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
red_text, blue_text = '\x1b[31m', '\x1b[34m'
white_background = '\x1b[107m'
msg = white_background + bright_style + red_text
msg += 'DeprecationWarning: This function will be deprecated in future. '
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
msg += reset_style
warnings.warn(msg)
register_op('one_hot', one_hot, '', opset)
register_op('im2col', im2col, '', opset)
register_op('topk', topk, '', opset)
register_op('softmax', softmax, '', opset)
register_op('constant_pad_nd', constant_pad_nd, '', opset)
register_op('reflection_pad1d', reflection_pad1d, '', opset)
register_op('reflection_pad2d', reflection_pad2d, '', opset)
register_op('reflection_pad3d', reflection_pad3d, '', opset)
register_op('avg_pool1d', avg_pool1d, '', opset)
register_op('avg_pool2d', avg_pool2d, '', opset)
register_op('avg_pool3d', avg_pool3d, '', opset)
register_op('adaptive_avg_pool1d', adaptive_avg_pool1d, '', opset)
register_op('adaptive_avg_pool2d', adaptive_avg_pool2d, '', opset)
register_op('adaptive_avg_pool3d', adaptive_avg_pool3d, '', opset)
register_op('masked_select', masked_select, '', opset)
register_op('upsample_nearest1d', upsample_nearest1d, '', opset)
register_op('upsample_nearest2d', upsample_nearest2d, '', opset)
register_op('upsample_nearest3d', upsample_nearest3d, '', opset)
register_op('upsample_linear1d', upsample_linear1d, '', opset)
register_op('upsample_bilinear2d', upsample_bilinear2d, '', opset)
register_op('upsample_trilinear3d', upsample_trilinear3d, '', opset)
register_op('upsample_bicubic2d', upsample_bicubic2d, '', opset)
register_op('new_full', new_full, '', opset)
register_op('grid_sampler', grid_sampler, '', opset)
register_op('cummax', cummax, '', opset)
register_op('cummin', cummin, '', opset)
register_op('roll', roll, '', opset)
interpolation_mode,
padding_mode,
align_corners=False):
'''
torch.nn.functional.grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=None)
Need convert interpolation_mode, padding_mode ? NO for simpler at now !!!
'''
return g.op('onnxservice::grid_sampler',
input,
grid,
interpolation_mode_i=interpolation_mode,
padding_mode_i=padding_mode,
align_corners_i=align_corners)
register_op('grid_sampler', grid_sampler, '', 11)
arguments_strModel = 'sintel-final'
SpyNet_model_dir = './models' # The directory of SpyNet's weights
import pdb
def normalize(input):
R = (input[:, 0:1, :, :] - 0.485) / 0.229
G = (input[:, 1:2, :, :] - 0.456) / 0.224
B = (input[:, 2:3, :, :] - 0.406) / 0.225
return torch.cat([R, G, B], 1)
def denormalize(input):
R = (input[:, 0:1, :, :] * 0.229) + 0.485
def register_extra_symbolics_for_openvino(opset=10):
assert opset >= 10
register_op('roi_feature_extractor', roi_feature_extractor_symbolics,
'mmdet_custom', opset)