Keras to TensorRT for TX2 (assertion error on Keras Shape operation)

I have been trying for sometime to get a custom model built and trained in Keras to compile to a TensorRT runtime engine for TX2. I began with a variational autoencoder with fully connected dense layers for the encoder and decoder. The commands and code for all of the steps is pasted below. I know this is alot but any help anyone could provide would be greatly appreciated.

commands:
host> ./train_vae.py <input_file> 64 (runs to completion and saves a _reconstruction_error_model.h5 file
host> ./keras_model_to_tf_pb.py _reconstruction_error_model.h5 reconst_error/BiasAdd ./ (runs to completion and saves a _reconstruction_error_model.pb file
host> scp .pb file to tx2

tx2> python scripts/convert_plan.py ./20180731_105005_reconstruction_error_model.pb ./20180731_105005_reconstruction_error_model.plan input 64 64 reconst_error/BiasAdd 1 0 float
Using output node reconst_error/BiasAdd
Converting to UFF graph
DEBUG: convert reshape to flatten node
DEBUG: convert reshape to flatten node
No. nodes: 39
UFF Output written to data/tmp.uff
UFFParser: parsing reconst_error/bias
UFFParser: parsing input
UFFParser: parsing encoder/flatten_1/Reshape
UFFParser: parsing dense_1/kernel
UFFParser: parsing encoder/dense_1/MatMul
UFFParser: parsing dense_1/bias
UFFParser: parsing encoder/dense_1/BiasAdd
UFFParser: parsing encoder/dense_1/Relu
UFFParser: parsing z_mean/kernel
UFFParser: parsing encoder/z_mean/MatMul
UFFParser: parsing z_mean/bias
UFFParser: parsing encoder/z_mean/BiasAdd
UFFParser: parsing dense_2/kernel
UFFParser: parsing decoder/dense_2/MatMul
UFFParser: parsing dense_2/bias
UFFParser: parsing decoder/dense_2/BiasAdd
UFFParser: parsing decoder/dense_2/Relu
UFFParser: parsing dense_3/kernel
UFFParser: parsing decoder/dense_3/MatMul
UFFParser: parsing dense_3/bias
UFFParser: parsing decoder/dense_3/BiasAdd
UFFParser: parsing decoder/dense_3/Sigmoid
UFFParser: parsing decoder/decoder_output/Shape
uff_to_plan: uff/UffParser.cpp:1370: std::shared_ptr UffParser::parseShape(const uff::Node&, const Fields&, NodesMap&): Assertion `order.size() == valuesOrdered.size()’ failed.

Keras Code to generate, train, and save the model:
########################################################################################################

train_vae.py

########################################################################################################

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from keras.layers import Lambda, Input, Dense, \
                         Flatten, Reshape, Subtract, Multiply
from keras.models import Model
from keras.losses import mse, binary_crossentropy
from keras.utils import plot_model
from keras import backend as K

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
import sys
import h5py
import datetime

import spectral_toolkit as st

tf_sess = tf.Session()
K.set_session(tf_sess)

spec_filename = sys.argv[1]
n_lookback = int(sys.argv[2])

base_filename = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + '_'

# reparameterization trick
# instead of sampling from Q(z|X), sample eps = N(0,I)
# z = z_mean + sqrt(var)*eps
def sampling(args):
    """Reparameterization trick by sampling fr an isotropic unit Gaussian.

    # Arguments:
        args (tensor): mean and log of variance of Q(z|X)

    # Returns:
        z (tensor): sampled latent vector
    """

    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    # by default, random_normal has mean=0 and std=1.0
    epsilon = K.random_normal(shape=(batch, dim))
    return z_mean + K.exp(0.5 * z_log_var) * epsilon

def error(args):
    ei = args
    sq = K.square(ei)
    s1 = K.sum(sq,0)
    s = K.sum(s1,0)
    return K.sqrt(s)

# Spectral dataset
n_infer = 1
f = h5py.File(spec_filename,'r')
d = f['spectrogram']
tt_idx = int(d.shape[0]*3/4)
d_train = d[:tt_idx,:]
d_test = d[tt_idx:,:]
x_train, y_train = st.create_dataset(d_train,n_lookback,n_infer)

time_len = x_train.shape[1]
freq_len = x_train.shape[2]
linear_len = time_len*freq_len

x_train = np.reshape(x_train,[-1, time_len, freq_len, 1])
x_train = st.normalize_tensor(x_train, 'uniform')
x_test = np.reshape(x_test,[-1, time_len, freq_len, 1])
x_test = st.normalize_tensor(x_test, 'uniform')

# network parameters
input_shape = (time_len, freq_len, 1)
intermediate_dim = 1024
batch_size = 128
latent_dim = 200
epochs = 100

# VAE model = encoder + decoder
# build encoder model
inputs = Input(shape=input_shape, name='input')
x = Flatten()(inputs)
x = Dense(intermediate_dim, activation='relu')(x)
z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)

# use reparameterization trick to push the sampling out as input
# note that "output_shape" isn't necessary with the TensorFlow backend
z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

# instantiate encoder model
encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')
encoder.summary()
plot_model(encoder,
           to_file=base_filename + 'encoder.png', show_shapes=True)

# build decoder model
latent_inputs = Input(shape=(latent_dim,), name='z_sampling')
x = Dense(intermediate_dim, activation='relu')(latent_inputs)
linear_outputs = Dense(linear_len, activation='sigmoid')(x)
outputs = Reshape(input_shape, name='decoder_output')(linear_outputs)

# instantiate decoder model
decoder = Model(latent_inputs, outputs, name='decoder')
decoder.summary()
plot_model(decoder,
            to_file=base_filename + 'decoder.png', show_shapes=True)

# build encoder_decoder model
reconst = decoder(encoder(inputs)[0])

# instantiate encoder_decoder model
encoder_decoder = Model(inputs, reconst, name='encoder_decoder')
encoder_decoder.summary()
plot_model(encoder_decoder,
            to_file='encoder_decoder.png',
            show_shapes=True)

# build reconstruction_error model
error_image = Subtract()([reconst, inputs])
#reconst_error = Lambda(error, output_shape=(1,), 
#                        name='reconst_error')(error_image)
error_vec = Flatten()(error_image)
square_error_vec = Multiply()([error_vec, error_vec])
reconst_error = Dense(1,kernel_initializer='ones',
                          bias_initializer='zeros',
                          trainable=False,
                          name='reconst_error',)(square_error_vec)

# instantiate reconstruction_error model
reconstruction_error = Model(inputs, reconst_error,
                                name='reconstruction_error')
reconstruction_error.summary()
plot_model(reconstruction_error,
            to_file='reconstruction_error.png',
            show_shapes=True)

# instantiate VAE model
outputs = decoder(encoder(inputs)[2])
vae = Model(inputs, outputs, name='vae_mlp')

# define loss function
reconstruction_loss = binary_crossentropy(K.flatten(inputs),
                                          K.flatten(outputs))

reconstruction_loss *= time_len * freq_len
kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
kl_loss = K.sum(kl_loss, axis=-1)
kl_loss *= -0.5
vae_loss = K.mean(reconstruction_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='adam')
vae.summary()
plot_model(vae,
           to_file=base_filename + 'vae.png',
           show_shapes=True)

# train the autoencoder
vae.fit(x_train,
        epochs=epochs,
        batch_size=batch_size,)

# save the models
encoder.save(base_filename + 'encoder_model.h5')
decoder.save(base_filename + 'decoder_model.h5')
encoder_decoder.save(base_filename + 'encoder_decoder_model.h5')
reconstruction_error.save(base_filename + 'reconstruction_error_model.h5')
vae.save(base_filename + 'vae_model.h5')

Code to “freeze” a keras model to a tf .pb

################################################################################################

keras_model_to_tf_pb.py

################################################################################################

#!/usr/bin/env python

import sys
import datetime

from keras.models import load_model
import keras.backend as K
from tensorflow.python.framework import graph_io
from tensorflow.python.tools import freeze_graph
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.training import saver as saver_lib

def convert_keras_to_pb(keras_model, out_names, models_dir, model_filename):
    model = load_model(keras_model)
    K.set_learning_phase(0)
    sess = K.get_session()
    saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2)
    checkpoint_path = saver.save(sess, './saved_ckpt', global_step=0,
                                    latest_filename='checkpoint_state')
    graph_io.write_graph(sess.graph, '.', 'tmp.pb')
    freeze_graph.freeze_graph('./tmp.pb', '',
                              False, checkpoint_path, out_names,
                              "save/restore_all", "save/Const:0",
                              models_dir+model_filename, False, "")

if __name__ == "__main__":
    keras_model = sys.argv[1]
    out_names = sys.argv[2]
    models_dir = sys.argv[3]
    model_filename = keras_model.replace('.h5','.pb')
    convert_keras_to_pb(keras_model, out_names, models_dir, model_filename)

Hi,

We try to reproduce this issue but meet some error in spectral_toolkit library.

ImportError: No module named spectral_toolkit

Guess that it is a custom module.
Could you share the implementation of spectral_toolkit with us?

Thanks.

yes sorry the code is below. The dataset can be any array of shape (n_samples, height, width, 1)

import numpy as np
import h5py

def create_dataset(spectrogram, n_lookback, n_infer):
    '''Returns a sliding window dataset from a spectrogram of shape (m,n)
       where m is the number of time instances, and n is the number of
       frequency bins.
       
       The dataset tuple is x,y where x is of shape (p,q,n) where:
         p is the number of samples,
         q is the number of lookback timesteps (n_lookback), and again,
         n is the number of frequency bins
       and y is of shape (p,r,n) where:
         r is the number of infer timesteps (n_infer).
'''

    x = []
    y = []
    for i in range(len(spectrogram)-n_lookback-n_infer):
        x.append(spectrogram[i:i+n_lookback,:])
        y.append(spectrogram[i+n_lookback:i+n_lookback+n_infer,:])
    return np.array(x), np.array(y)

def normalize_tensor(t, norm_type='uniform'):
    if norm_type == 'uniform':
        return (t-np.min(t))/(np.max(t)-np.min(t))
    elif norm_type == 'normal':
        return (t-np.mean(t))/np.std(t)

Hi,

We found you have open another Keres topic:
https://devtalk.nvidia.com/default/topic/1038119

Is this original issue solved?
Thanks

No this isnt solved yet. i tried a simpler model and ran into that number of weights issue.

Hi,

Could you try if the converter shared in topic 1038119 also fixes this issue?
https://devtalk.nvidia.com/default/topic/1038119/jetson-tx2/incorrect-number-of-kernel-weights-for-keras-to-pb-model-expecting-3x-as-many-as-available-/post/5275065/#5275065

Thanks.