Source code for mxnet.contrib.tensorrt

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

""" Module to enable the use of TensorRT optimized graphs."""
import os

[docs]def set_use_fp16(status): """ Set an environment variable which will enable or disable the use of FP16 precision in TensorRT Note: The mode FP16 force the whole TRT node to be executed in FP16 :param status: Boolean, True if TensorRT should run in FP16, False for FP32 """ os.environ["MXNET_TENSORRT_USE_FP16"] = str(int(status))
[docs]def get_use_fp16(): """ Get an environment variable which describes if TensorRT is currently running in FP16 :return: Boolean, true if TensorRT is running in FP16, False for FP32 """ return bool(int(os.environ.get("MXNET_TENSORRT_USE_FP16", 1)) == 1)
[docs]def init_tensorrt_params(sym, arg_params, aux_params): """ Set weights in attributes of TensorRT nodes :param sym: Symbol, the symbol graph should contains some TensorRT nodes :param arg_params: arg_params :param aux_params: aux_params :return arg_params, aux_params: remaining params that are not in TensorRT nodes """ arg_params = arg_params.copy() aux_params = aux_params.copy() for s in sym.get_internals(): new_params_names = "" tensorrt_params = {} if 'subgraph_params_names' in s.list_attr(): keys = s.list_attr()['subgraph_params_names'].split(';') for k in keys: if k in arg_params: new_params_names += k + ";" tensorrt_params['subgraph_param_' + k] = arg_params[k] arg_params.pop(k) elif k in aux_params: new_params_names += k + ";" tensorrt_params['subgraph_param_' + k] = aux_params[k] aux_params.pop(k) new_attrs = {} for k, v in tensorrt_params.items(): new_attrs[k] = str(v.handle.value) if len(new_attrs) > 0: s._set_attr(**new_attrs) s._set_attr(subgraph_params_names=new_params_names[:-1]) return arg_params, aux_params