Ошибка TensorRT Engine из-за ошибки парсера UFF

Я пытаюсь преобразовать мою замороженную модель Tensorflow на основе архитектуры pix2pix в движок NVIDIA TensorRT. У меня версия TensorRT 3. Хотя мои слои поддерживаются TensorRT 3, некоторые из них не поддерживаются парсером UFF. Проблема, с которой я столкнулся, связана с строкой:

mean, variance = tf.nn.moments(input, axes=[0, 1, 2], keep_dims=False)

Одна из внутренних операций - это squared difference, которая не поддерживается TensorRT, при выполнении возвращающая следующую ошибку:

Преобразование в график UFF Предупреждение: keep_dims не поддерживается, игнорируется ...

Предупреждение: функция преобразования для слоя: SquaredDifference еще не зарегистрирована.

Преобразование в пользовательский оператор SquaredDifference>

generate_output / generator / decoder_2 / batchnorm / Moments / SquaredDifference имя:

"generate_output / generator / decoder_2 / batchnorm / Moments / SquaredDifference"

op: "SquaredDifference"

ввод: "generate_output / generator / decoder_2 / batchnorm / Identity"

ввод: "generate_output / generator / decoder_2 / batchnorm / Moments / StopGradient"

attr {

ключ: «Т»

ценить {

тип: DT_FLOAT

}

}

Отслеживание (последний вызов последний):

Файл "scripts / convert_plan.py", строка 71, в data_type

Файл "scripts / convert_plan.py", строка 22, в файле frozenToPlan text = False,

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/conversion_helpers.py", строка 103, в from_tensorflow_frozen_model return from_tensorflow (graphdef, output_nodes, ** kwargs)

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/conversion_helpers.py", строка 75, в from_tensorflow name = "main")

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 64, в convert_tf2uff_graph uff_graph, input_replacements)

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 51, в convert_tf2uff_node op, name, tf_node, inputs, uff_graph, tf_nodes = tf_nodes)

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 28, в convert_layer fields = cls.parse_tf_attrs (tf_node.attr)

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 177, в parse_tf_attrs для ключа, val в attrs.items ()}

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 177, для ключа, val в attrs.items ()}

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 172, в parse_tf_attr_value return cls.convert_tf2uff_field (code, val)

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 146, в convert_tf2uff_field возвращает TensorFlowToUFFConverter.convert_tf2numpy_dtype (val)

Файл "/usr/local/lib/python2.7/dist-packages/uff/converters/tensorflow/converter.py", строка 74, в convert_tf2numpy_dtype return np.dtype (dt [dtype])

TypeError: индексы списка должны быть целыми числами, а не AttrValue

Я попытался заменить tf.nn.moments двумя альтернативными методами:

1. mean = tf.reduce_mean(input, axis = [0, 1, 2], keep_dims=False)

   devs_squared = tf.multiply(input - mean, input - mean)

   variance = tf.reduce_mean(devs_squared, axis = [0, 1, 2], keep_dims=False)

Хотя все слои поддерживаются TensorRT 3, reduce_mean не поддерживается парсером UFF и вызывает следующую ошибку:

UFFParser: Ошибка парсера:

generate_output / generator / encoder_2 / batchnorm / Mean: оператор сокращения не поддерживается

2. Второй способ:

   mean = tf.keras.backend.sum(input, axis = [0, 1, 2])

   div_factor_1 = input.get_shape().as_list()[0]*input.get_shape().as_list()[1]*input.get_shape().as_list()[2]

   mean = mean * (1 / div_factor_1)

   devs_squared = tf.multiply(input - mean, input - mean)

   variance = tf.keras.backend.sum(devs_squared, axis = [0, 1, 2])

   div_factor_2 = devs_squared.get_shape().as_list()[0]*devs_squared.get_shape().as_list()[1]*devs_squared.get_shape().as_list()[2]

   variance = variance * (1 / div_factor_2)

Но я снова сталкиваюсь с ошибкой парсера UFF:

UFFParser: ошибка парсера: generate_output / generator / encoder_2 / batchnorm / Sum: оператор сокращения не поддерживается

Я также использовал reduce_sum в альтернативе 2 вместо keras.backend.sum, чтобы встретить ту же ошибку.