Tensorflow 1.11.0
TensorRT 5.0.0.10
Ubuntun 16.04
Python 2.7
Cuda 9.0
Cudnn 7.1.2
I met a problem:
[TensorRT] ERROR: Internal error: could not find any implementation for node 2-layer MLP, try increasing the workspace size with IBuilder::setMaxWorkspaceSize()
[TensorRT] ERROR: ../builder/tacticOptimizer.cpp (1228) - OutOfMemory Error in computeCosts: 0
I am using the sample form the tutorial which located in “TensorRT-5.0.0.10/python/sample/end_to_end_mnist”.
This sample works fine when I simply run it in terminal.
But when I use the model(the keras model in “end_to_end_mnist”) and try to build an engine from scratch
it get an error.
I follow Overview — NVIDIA TensorRT Documentation have frozen the model and parse it successfully but fail to build an engine
I try both decreasing the batchsize and increasing the workspace but they did not work.
Does anyone know how to solve this problem?
FILE_ONE:
import tensorflow as tf
import numpy as np
import tensorrt as trt
import sys,os
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
print(type(x_train))
NUM_TRAIN = 60000
NUM_TEST = 10000
x_train = np.reshape(x_train, (NUM_TRAIN, 28, 28, 1))
x_test = np.reshape(x_test, (NUM_TEST, 28, 28, 1))
model = tf.keras.models.Sequential()
model.add(tf.keras.layers.InputLayer(input_shape=[28, 28, 1]))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(512, activation=tf.nn.relu))
model.add(tf.keras.layers.Dense(10, activation=tf.nn.softmax))
model.compile(optimizer=‘adam’, loss=‘sparse_categorical_crossentropy’, metrics=[‘accuracy’])
Train the model on the data
model.fit(x_train, y_train, epochs=1, verbose=1)
Evaluate the model on test data
model.evaluate(x_test, y_test)
output_names = model.output.op.name
sess = tf.keras.backend.get_session()
frozen_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), [output_names])
frozen_graph = tf.graph_util.remove_training_nodes(frozen_graph)
Save the model
with open(“testkeras_uff”, “wb”) as ofile:
ofile.write(frozen_graph.SerializeToString())
class ModelData(object):
MODEL_FILE = os.path.join(os.path.dirname(file), “testkeras_uff”)
INPUT_NAME =“input_1”
INPUT_SHAPE = (1, 28, 28)
OUTPUT_NAME = “dense_1/Softmax”
BEFORE RUNING FILE TWO U SHOULD “convert-to-uff testkeras_uff” in terminal
FILE_TWO:
import tensorflow as tf
import numpy as np
import tensorrt as trt
import sys,os
import pycuda.driver as cuda
import pycuda.autoinit
import argparse
model_file=“testkeras_uff.uff”
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
# Parse the Uff Network
parser.register_input(“input_1”, (1,28,28))
parser.register_output(“dense_1/Softmax”)
parser.parse(model_file, network)
# Build and return an engine.
builder.build_cuda_engine(network)
builder.max_batch_size = 1
builder.max_workspace_size = 1 <<20
with builder.build_cuda_engine(network) as engine:
h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=np.float32)
# Allocate device memory for inputs and outputs.
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
# Create a stream in which to copy inputs/outputs and run inference.
stream=cuda.Stream()
with engine.create_execution_context() as context:
# Transfer input data to the GPU.
cuda.memcpy_htod_async(d_input, h_input, stream)
# Run inference.
context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
# Transfer predictions back from the GPU.
cuda.memcpy_dtoh_async(h_output, d_output, stream)
# Synchronize the stream
stream.synchronize()
# Return the host output.
Thanks.