ChatBot: Use TensorRT to Inference a TensorFlow Model

Demonstrate how to use TensorRT to accelerate TensorFlow inference on Jetson.




You can learn following things with this sample:

1. Convert an LSTM TensorFlow model into UFF format.
2. Run a UFF model on Jetson with Python (includes a C++ to Python TensorRT wrapper).

We also include chatbot training and x86 chatbot source(pure python) for user reference.

---

Environment

Device

1. LINK Jetson TX2
2. LINK JetPack 3.2

Please flash your device with JetPack3.2 first.

sudo apt-get install python-pip
sudo apt-get install swig
sudo pip install numpy

Host

1. LINK Ubuntu16.04
2. LINK CUDA Toolkit 9
3. LINK TensorRT 3

Please get CUDA 9 and TensorRT 3 installed first.

sudo apt-get install python-pip
sudo pip install tensorflow
sudo apt-get install swig
sudo pip install numpy

Quick Try

We also attach a ChatBot model for user reference.

Notes

1. Don't blame him. He learns to talk with a small database(Cornell).
2. This model is trained on ONE Tesla P100 GPU and takes around 15 days to finish.
3. You can have the training job done much faster with a DGX-1 server.
4. Word vector is 4096. A word not in our vocabulary is replaced by '_'.
5. Our ChatBot is sensitive to the symbol. Please don't forget your symbol at the end of the sentence.
6. Our model is GAN.

Execution steps

Host: convert TF model into UFF format

$ git clone https://github.com/AastaNV/ChatBot.git
$ cd $CHATBOT_ROOT
$ python src/tf_to_uff/tf_to_trt.py model/ID210_649999 model/ID210_649999.uff
$ scp model/ID210_649999.uff nvidia@[device IP]:$CHATBOT_ROOT/model/

Device: create TensorRT engine with the converted UFF file

$ 
$ git clone https://github.com/AastaNV/ChatBot.git
$ wget https://raw.githubusercontent.com/numpy/numpy/master/tools/swig/numpy.i -P $CHATBOT_ROOT/src/
$ cd $CHATBOT_ROOT
$ make
$ python chatbot.py model/ID210.pickle model/ID210_649999.uff 

Examples

Hi, ChatBot!
hello .

How are you today?
fine .

How about to go out for dinner?
okay .

What do you want to eat?
i don't know .

Steak?
well , i don't think so .

Sandwich?
what ?

The coffee shop on the corner.
yes .

Great!
i was just wondering .

Wonder what?
i don't know .

Okay. YOu are a artist, right?
oh , no . . .

Anyway, see you tonight.
what ?

See you...
i know .

Okay... bye bye.
goodnight .

Export TensorFlow Model to UFF Format

Please run this step on host

Prepare following files from TensorFlow:

• [myModel].data-00000-of-00001
• [myModel].index
• [myModel].meta

$ cd $CHATBOT_ROOT/src/tf_to_uff
$ python tf_to_trt.py [myModel] [myModel].uff

Run UFF Model on Jetson

Please run this step on device

Prepare following files in model directory:

• [database].pickle
• [myModel].uff

$ wget https://raw.githubusercontent.com/numpy/numpy/master/tools/swig/numpy.i -P $CHATBOT_ROOT/src/
$ cd $CHATBOT_ROOT/
$ make
$ python chatbot.py model/[database].pickle model/[myModel].uff 

Training

Please run this step on host

We also share our training code for user reference.



Prepare following files from Cornell Movie Dialogs Corpus:

• movie_conversations.txt
• movie_lines.txt

$ cd $CHATBOT_ROOT/src/training
$ python parser_for_cornell.py
$ python main.py 

Run ChatBot on X86-based Linux Machine

Please run this step on host

We also provide source to run TensorRT on x86-based Machine
Prepare following files:

• [myModel].data-00000-of-00001
• [myModel].index
• [myModel].meta
• [database].pickle

$ cd $CHATBOT_ROOT/src/tf_to_uff
$ python tf_to_trt.py [myModel] [myModel].uff
$ python chatbot.py [database].pickle [myModel].uff