Let’s say I define a very minimal network and have its frozen graph as below.
""" minial example of exporting network to frozen graph """
with tf.Graph().as_default():
""" some network """
image = tf.placeholder(tf.float32, [1, cfg.IMAGE_HEIGHT, cfg.IMAGE_WIDTH, 3])
net = slim.conv2d(image, 4, [3, 3], scope='conv1')
net = slim.max_pool2d(net, [64, 64], scope='pool1')
net = tf.reshape(net, (-1, 2), name='reshape1')
net = tf.nn.softmax(net, name='softmax')
gpu_config = tf.ConfigProto(allow_soft_placement=True)
gpu_config.gpu_options.allow_growth = True
with tf.Session(config=gpu_config) as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
""" specify tensors that I will need when doing inference """
output_tensors = [net]
output_names = [n.op.name for n in output_tensors]
print("output tensors: ", output_names)
graphdef = tf.get_default_graph().as_graph_def()
frozen_graph = tf.graph_util.convert_variables_to_constants(sess, graphdef, output_names)
frozen_graph = tf.graph_util.remove_training_nodes(frozen_graph)
tf.train.write_graph(frozen_graph, '.', 'reshape-test.pb', as_text=False)
I convert it to UFF as below,
$ convert-to-uff tensorflow --input-file reshape-test.pb -O rpn_process -p reshape-test-uff-custom.py
where the “reshape-test-uff-custom.py” is:
import graphsurgeon as gs
import tensorflow as tf
reshape = gs.create_node("rpn_process", op="rpn_process", dtype=tf.float32)
namespace_plugin_map = {"reshape1": reshape}
def preprocess(dynamic_graph):
dynamic_graph.collapse_namespaces(namespace_plugin_map)
It simply replaces “reshape1” with a custom plugin layer “rpn_process”.
And then, I implemented the “rpn_process” custom layer, and overrode “PluginFactory::isPlugin” as below, so that the layer name “_rpn_process” is known.
class PluginFactory : public nvinfer1::IPluginFactory, public nvuffparser::IPluginFactory
{
public:
IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override
{
std::cout << layerName << std::endl;
assert(isPlugin(layerName));
......
}
......
bool isPlugin(const char* name) override
{
return !strcmp(name, "_rpn_process");
}
}
However, “assert(isPlugin(layerName));” results in an assertion error during runtime, since internally the layer name changed to “_rpn_process_HL_1804289383”.
Why is there a strange postfix in my custom layer? It happens when I replace some part of a given graph with a custom plugin layer, and declare the plugin layer as an output node.
Could you tell me what I am doing wrong?