EFanZh
commented
4 years ago I think I am having the same problem converting a simple MNIST CNN model. Here is my code for reproducing (You need to install tensorflow and tf2onnx to run this):
from multiprocessing import Pool
import numpy
import tensorflow as tf
from tensorrt import Builder, BuilderFlag, IInt8EntropyCalibrator2, Logger, NetworkDefinitionCreationFlag, OnnxParser
(_TRAIN_IMAGES, _TRAIN_LABELS), _ = tf.keras.datasets.mnist.load_data()
_TRAIN_IMAGES = numpy.expand_dims(a=_TRAIN_IMAGES, axis=-1).astype(numpy.float32)
_TRAIN_LABELS = tf.keras.utils.to_categorical(y=_TRAIN_LABELS, num_classes=10)
class _Calibrator(IInt8EntropyCalibrator2):
def __init__(self):
super().__init__()
self._batch_size = 1
self._cache = None
def get_batch(self, names, p_str=None):
raise NotImplementedError
def get_batch_size(self):
return self._batch_size
def read_calibration_cache(self):
return self._cache
def write_calibration_cache(self, cache):
self._cache = cache
def _assert(value):
if not value:
raise AssertionError
def _get_frozen_graph_model():
with tf.compat.v1.Session(graph=tf.Graph()) as session:
model = tf.keras.Sequential(layers=[
tf.keras.layers.Conv2D(filters=32, kernel_size=[3, 3], activation='relu', input_shape=[28, 28, 1]),
tf.keras.layers.Conv2D(filters=64, kernel_size=[3, 3], activation='relu'),
tf.keras.layers.MaxPooling2D(pool_size=[2, 2]),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(units=128, activation='relu'),
tf.keras.layers.Dense(units=10, activation='softmax'),
])
model.compile(optimizer=tf.keras.optimizers.SGD(), loss='binary_crossentropy', metrics=['accuracy'])
model.fit(x=_TRAIN_IMAGES, y=_TRAIN_LABELS)
return (tf.compat.v1.graph_util.convert_variables_to_constants(sess=session,
input_graph_def=session.graph_def,
output_node_names=[model.output.op.name]),
model.input.name,
model.output.name)
def _tf_to_onnx(graph_def, input_name, output_name):
from onnx import defs
from tf2onnx import tfonnx
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def=graph_def, name='')
onnx_model = tfonnx.process_tf_graph(tf_graph=graph,
opset=defs.onnx_opset_version(),
input_names=[input_name],
output_names=[output_name])
return onnx_model.make_model('').SerializeToString()
def _onnx_to_tensorrt(onnx_model):
batch_size = 1
with Logger() as logger, \
Builder(logger) as builder, \
builder.create_network(1 << int(NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
OnnxParser(network, logger) as onnx_parser:
_assert(onnx_parser.parse(onnx_model))
builder.max_batch_size = batch_size
builder_config = builder.create_builder_config()
optimization_profile = builder.create_optimization_profile()
for i in range(network.num_inputs):
input_tensor = network.get_input(i)
shape = (batch_size,) + input_tensor.shape[1:]
optimization_profile.set_shape(input=input_tensor.name, min=shape, opt=shape, max=shape)
builder_config.add_optimization_profile(optimization_profile)
builder_config.set_flag(BuilderFlag.INT8)
builder_config.int8_calibrator = _Calibrator()
cuda_engine = builder.build_engine(network, builder_config)
_assert(cuda_engine)
return cuda_engine
def _get_onnx_model():
graph_def, input_name, output_name = _get_frozen_graph_model()
return _tf_to_onnx(graph_def=graph_def, input_name=input_name, output_name=output_name)
def main():
with Pool(processes=1) as pool:
# Run in another process to make sure GPU memory used by TensorFlow gets freed.
onnx_model = pool.apply(_get_onnx_model)
_onnx_to_tensorrt(onnx_model=onnx_model)
if __name__ == '__main__':
main()
rmccorm4
wyp19960713
wanghr323
Description
[TensorRT] VERBOSE: Engine generation completed in 3.27319 seconds. [TensorRT] VERBOSE: Calculating Maxima [TensorRT] ERROR: ../rtSafe/safeRuntime.cpp (25) - Cuda Error in allocate: 2 (out of memory) [TensorRT] ERROR: ../rtSafe/safeRuntime.cpp (25) - Cuda Error in allocate: 2 (out of memory) When the calibrator is activated,these errors occur. The following is my code of calibrator:
Environment
TensorRT Version: 7.0.0.11 GPU Type: RTX 2070 Nvidia Driver Version: 440.82 CUDA Version: 10.0 CUDNN Version: 7.6.4 Operating System + Version:Ubuntu16.04 Python Version (if applicable): 3.7.6 TensorFlow Version (if applicable): 1.14 PyTorch Version (if applicable): Baremetal or Container (if container which image + tag):
Relevant Files
Steps To Reproduce