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)