diff --git a/tftrt/examples/cpp/image-classification/CMakeLists.txt b/tftrt/examples/cpp/image-classification/CMakeLists.txt new file mode 100644 index 000000000..c7cad6e43 --- /dev/null +++ b/tftrt/examples/cpp/image-classification/CMakeLists.txt @@ -0,0 +1,47 @@ +cmake_minimum_required(VERSION 3.13) +project(TF_TRT_Example) + +#------------------------------------------------------------- +# Configuration +#------------------------------------------------------------- +set(CMAKE_CXX_STANDARD 14) +set(TF_INSTALL_DIR "/usr/local" CACHE PATH "Path to Tensorflow install directory") + +#------------------------------------------------------------- +# Dependencies +#------------------------------------------------------------- +find_library(tf_shared_lib2 NAMES libtensorflow_cc.so.2 HINTS ${TF_INSTALL_DIR}/lib/tensorflow) +find_library(tf_framework_shared_lib2 NAMES libtensorflow_framework.so.2 HINTS ${TF_INSTALL_DIR}/lib/python3.8/dist-packages/tensorflow) +find_path(trt_include_path NAME NvInfer.h HINTS) +get_filename_component(tf_dir ${tf_shared_lib2} DIRECTORY) +get_filename_component(tf_python_dir ${tf_framework_shared_lib2} DIRECTORY) + +set(tf_shared_lib ${tf_dir}/libtensorflow_cc.so) +set(tf_framework_shared_lib ${tf_python_dir}/libtensorflow_framework.so) + +add_custom_command(OUTPUT ${tf_framework_shared_lib} ${tf_shared_lib} + COMMAND ln -s ${tf_shared_lib2} ${tf_shared_lib} + COMMAND ln -s ${tf_framework_shared_lib2} ${tf_framework_shared_lib} + COMMENT "Generating legacy symbolic link") + +add_custom_target(tf_symlinks DEPENDS ${tf_framework_shared_lib} ${tf_shared_lib}) + +#----------------------------------------------------------- +# Example Targets +#----------------------------------------------------------- +add_executable(tf_trt_example main.cc mnist.h mnist.cc) + +target_link_libraries(tf_trt_example tensorflow_cc) +target_link_libraries(tf_trt_example tensorflow_framework) + +target_compile_options(tf_trt_example PRIVATE -D_GLIBCXX_USE_CXX11_ABI=0 -DGOOGLE_CUDA -DGOOGLE_TENSORRT) + +target_link_directories(tf_trt_example PRIVATE ${tf_python_dir}) +target_link_directories(tf_trt_example PRIVATE ${tf_dir}) + +target_compile_options(tf_trt_example PRIVATE -Wl,-rpath=${tf_python_dir}) + +target_include_directories(tf_trt_example PRIVATE ${tf_python_dir}/include) +target_include_directories(tf_trt_example PRIVATE ${trt_include_path}) + +add_dependencies(tf_trt_example tf_symlinks) diff --git a/tftrt/examples/cpp/image-classification/README.md b/tftrt/examples/cpp/image-classification/README.md new file mode 100644 index 000000000..72ed5b76c --- /dev/null +++ b/tftrt/examples/cpp/image-classification/README.md @@ -0,0 +1,49 @@ +# TF-TRT example for conversion in C++ + +## Introduction + +This directory contains example code to demonstrate TF-TRT conversion using the C++ API. + +### Acknowledgment + +The MNIST inference example is based on https://github.com/bmzhao/saved-model-example + +## How to run + +### Build TF +``` +git clone https://github.com/tensorflow/tensorrt.git +git clone https://github.com/tensorflow/tensorflow.git tensorflow-source +mkdir bazel-cache +docker run --gpus=all --rm -it --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 -v $PWD:/workspace -w /workspace -v $PWD/bazel-cache:/root/.cache/bazel nvcr.io/nvidia/tensorflow:21.06-tf2-py3 + +# Inside the container +cp /opt/tensorflow/nvbuild* /opt/tensorflow/bazel_build.sh . +./nvbuild.sh --noclean --v2 +``` + +### Build the TF-TRT example +``` +cd tensorrt/tftrt/examples/cpp/image-classification +mkdir build && cd build +cmake .. +make +``` + +### Train the model and save it +``` +python mnist_train.py +``` + +### Run TF-TRT conversion and infer the converted model +Get input data +``` +cd /workspace/tensorrt/tftrt/examples/cpp/image-classification +wget -O - http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz | gunzip > t10k-images.idx3-ubyte + +``` +Run inference +``` +cd /workspace/tensorrt/tftrt/examples/cpp/image-classification/build +TF_CPP_VMODULE=trt_convert=2,trt_optimization_pass=2,trt_engine_utils=2,trt_engine_op=2,segment=2,trt_shape_optimization_profiles=2,trt_lru_cache=2,convert_graph=2,trt_engine_resource_ops=2 ./tf_trt_example --saved_model_dir=/workspace/tensorflow-source/tf_trt_cpp_example/mnist_model --mnist_data=/workspace/tensorflow-source/tf_trt_cpp_example/t10k-images.idx3-ubyte +``` diff --git a/tftrt/examples/cpp/image-classification/main.cc b/tftrt/examples/cpp/image-classification/main.cc new file mode 100644 index 000000000..3b169c870 --- /dev/null +++ b/tftrt/examples/cpp/image-classification/main.cc @@ -0,0 +1,206 @@ +/* Copyright 2021 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include + +#include "absl/strings/string_view.h" +#include "mnist.h" +#include "tensorflow/cc/saved_model/loader.h" +#include "tensorflow/compiler/tf2tensorrt/trt_convert_api.h" +#include "tensorflow/core/framework/tensor.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/public/session.h" +#include "tensorflow/core/util/command_line_flags.h" + +// Returns the name of nodes listed in the signature definition. +std::vector +GetNodeNames(const google::protobuf::Map + &signature) { + std::vector names; + for (auto const &item : signature) { + absl::string_view name = item.second.name(); + // Remove tensor suffix like ":0". + size_t last_colon = name.find_last_of(':'); + if (last_colon != absl::string_view::npos) { + name.remove_suffix(name.size() - last_colon); + } + names.push_back(std::string(name)); + } + return names; +} + +// Loads a SavedModel from export_dir into the SavedModelBundle. +tensorflow::Status LoadModel(const std::string &export_dir, + tensorflow::SavedModelBundle *bundle, + std::vector *input_names, + std::vector *output_names) { + + tensorflow::RunOptions run_options; + TF_RETURN_IF_ERROR(tensorflow::LoadSavedModel(tensorflow::SessionOptions(), + run_options, export_dir, + {"serve"}, bundle)); + + // Print the signature defs. + auto signature_map = bundle->GetSignatures(); + for (const auto &name_and_signature_def : signature_map) { + const auto &name = name_and_signature_def.first; + const auto &signature_def = name_and_signature_def.second; + std::cerr << "Name: " << name << std::endl; + std::cerr << "SignatureDef: " << signature_def.DebugString() << std::endl; + } + + // Extract input and output tensor names from the signature def. + const tensorflow::SignatureDef &signature = signature_map["serving_default"]; + *input_names = GetNodeNames(signature.inputs()); + *output_names = GetNodeNames(signature.outputs()); + + return tensorflow::Status::OK(); +} + +// We prepare to input sets, with batch size 1 and batch size 4. +tensorflow::Status +LoadInputs(const std::string &mnist_path, + std::vector> *inputs) { + // Load the MNIST images from the given path. + mnist::MNISTImageReader image_reader(mnist_path); + TF_RETURN_IF_ERROR(image_reader.ReadMnistImages()); + + // Convert the first image to a tensorflow::Tensor. + tensorflow::Tensor input_image = image_reader.MNISTImageToTensor(0, 1); + std::cout << "Created input tensor with shape " + << input_image.shape().DebugString() << std::endl; + inputs->push_back({input_image}); + + // Create another tensor with batch size 4. + input_image = image_reader.MNISTImageToTensor(0, 4); + std::cout << "Created input tensor with shape " + << input_image.shape().DebugString() << std::endl; + inputs->push_back({input_image}); + + for (int i = 0; i < 4; i++) { + std::cout << "Input image " << i << std::endl; + std::cout << image_reader.images[i] << std::endl; + } + return tensorflow::Status::OK(); +} + +#define TFTRT_ENSURE_OK(x) \ + do { \ + tensorflow::Status s = x; \ + if (!s.ok()) { \ + std::cerr << __FILE__ << ":" << __LINE__ << " " << s.error_message() \ + << std::endl; \ + return 1; \ + } \ + } while (0) + +int main(int argc, char **argv) { + // One can use a command line arg to specify the saved model directory. + // Note that the model has to be trained with input size [None, 28, 28]. + // Currently, it is assumed that the model is already frozen (variables are + // converted to constants). + std::string export_dir = + "/workspace/tensorflow-source/tf_trt_cpp_example/mnist_model_frozen"; + std::string mnist_data_path = + "/workspace/tensorflow-source/tf_trt_cpp_example/" + "t10k-images.idx3-ubyte"; + bool frozen_graph = false; + std::vector flag_list = { + tensorflow::Flag("saved_model_dir", &export_dir, + "Path to saved model directory"), + tensorflow::Flag("mnist_data", &mnist_data_path, "Path to MNIST images"), + tensorflow::Flag("frozen_graph", &frozen_graph, + "Assume graph is frozen and use TF-TRT API for frozen " + "graphs")}; + const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list); + + if (!parse_result) { + std::string usage = tensorflow::Flags::Usage(argv[0], flag_list); + LOG(ERROR) << usage; + return -1; + } + tensorflow::port::InitMain(argv[0], &argc, &argv); + + tensorflow::SavedModelBundle bundle; + std::vector input_names; + std::vector output_names; + + // Load the saved model from the provided path. + TFTRT_ENSURE_OK(LoadModel(export_dir, &bundle, &input_names, &output_names)); + + // Prepare input tensors + std::vector> inputs; + TFTRT_ENSURE_OK(LoadInputs(mnist_data_path, &inputs)); + + // Run TF-TRT conversion + tensorflow::tensorrt::TfTrtConversionParams params; + params.use_dynamic_shape = true; + params.profile_strategy = tensorflow::tensorrt::ProfileStrategy::kOptimal; + tensorflow::StatusOr status_or_gdef; + if (frozen_graph) { + status_or_gdef = tensorflow::tensorrt::ConvertAndBuild( + bundle.meta_graph_def.graph_def(), input_names, output_names, inputs, + params); + } else { + status_or_gdef = tensorflow::tensorrt::ConvertAndBuild( + &bundle, "serving_default", inputs, params); + } + if (!status_or_gdef.ok()) { + std::cerr << "Error converting the graph" << status_or_gdef.status() + << std::endl; + return 1; + } + tensorflow::GraphDef &converted_graph_def = status_or_gdef.ValueOrDie(); + tensorflow::Session *session = nullptr; + TFTRT_ENSURE_OK(NewSession(tensorflow::SessionOptions(), &session)); + bundle.session.reset(session); + TFTRT_ENSURE_OK(bundle.session->Create(converted_graph_def)); + + // Infer the converted model + for (auto const &input : inputs) { + std::vector> input_pairs; + for (int i = 0; i < input_names.size(); i++) { + input_pairs.push_back({input_names.at(i), input.at(i)}); + } + std::vector output_tensors; + for (int i = 0; i < output_names.size(); i++) { + output_tensors.push_back({}); + } + std::cout << "Inferring the model" << std::endl; + TFTRT_ENSURE_OK( + session->Run(input_pairs, output_names, {}, &output_tensors)); + + for (const auto &output_tensor : output_tensors) { + const auto &vec = output_tensor.flat_inner_dims(); + int m = output_tensor.dim_size(1); + for (int k = 0; k < output_tensor.dim_size(0); k++) { + float max = 0; + int argmax = 0; + for (int i = 0; i < m; ++i) { + if (vec(i + k * m) > max) { + argmax = i; + } + } + std::cerr << "Sample " << k << ", Predicted Number: " << argmax + << std::endl; + } + } + } +} diff --git a/tftrt/examples/cpp/image-classification/mnist.cc b/tftrt/examples/cpp/image-classification/mnist.cc new file mode 100644 index 000000000..21d0d6c26 --- /dev/null +++ b/tftrt/examples/cpp/image-classification/mnist.cc @@ -0,0 +1,131 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "mnist.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include +#include +#include +#include +#include +#include + +namespace mnist { + +namespace { + +// This function interprets a 4 byte array as an unsigned, +// big-endian integer. sizeof(data) must be 4. +// From Rob Pike's Byte Order Fallacy: +// https://commandcenter.blogspot.com/2012/04/byte-order-fallacy.html +uint32_t ConvertBigEndian(unsigned char data[]) { + uint32_t result = + (data[3] << 0) | (data[2] << 8) | (data[1] << 16) | (data[0] << 24); + return result; +} + +} // namespace + +std::ostream &operator<<(std::ostream &os, const MNISTImage &image) { + for (int row = 0; row < MNISTImage::kSize; ++row) { + for (int column = 0; column < MNISTImage::kSize; ++column) { + os << (image.buf[column + MNISTImage::kSize * row] > 0 ? "X " : " "); + } + os << "\n"; + } + return os; +} + +tensorflow::Tensor MNISTImageReader::MNISTImageToTensor(int offset, + int batch_size) { + // https://github.com/tensorflow/tensorflow/issues/8033#issuecomment-520977062 + tensorflow::Tensor input_image( + tensorflow::DT_FLOAT, + tensorflow::TensorShape( + {batch_size, MNISTImage::kSize, MNISTImage::kSize})); + float *img_tensor_flat = input_image.flat().data(); + constexpr int N = MNISTImage::kSize * MNISTImage::kSize; + for (int i = 0; i < batch_size; i++) { + std::copy_n(images[i].buf.data() + offset * N, N, img_tensor_flat + i * N); + } + return input_image; +} + +MNISTImageReader::MNISTImageReader(const std::string &path) + : mnist_path_(path) {} + +// MNIST's file format is documented here: http://yann.lecun.com/exdb/mnist/ +// Note(bmzhao): The serialized integers are in big endian, and the magic +// number is documented to be 2051. +tensorflow::Status MNISTImageReader::ReadMnistImages() { + std::ifstream image_file(mnist_path_, std::ios::binary); + if (!image_file.is_open()) { + return tensorflow::errors::NotFound("Error opening file", mnist_path_, ": ", + std::strerror(errno)); + } + + uint32_t num_images; + uint32_t num_rows; + uint32_t num_columns; + + auto ReadBigEndian = [](std::ifstream &image_file) { + unsigned char buf[4]; + image_file.read(reinterpret_cast(&buf[0]), sizeof(buf)); + return ConvertBigEndian(buf); + }; + + // Read the magic number. + uint32_t magic_number = ReadBigEndian(image_file); + if (magic_number != 2051) { + return tensorflow::errors::Internal("Magic Number of Mnist Data File ", + mnist_path_, " was ", magic_number, + " expected 2051"); + } + + // Read the number of images. + num_images = ReadBigEndian(image_file); + std::cout << "Number of images " << num_images << std::endl; + // Read the number of rows. + num_rows = ReadBigEndian(image_file); + if (num_rows != MNISTImage::kSize) { + return tensorflow::errors::FailedPrecondition( + "Num Rows of Mnist Data File was ", num_rows, " expected 28"); + } + + // Read the number of columns + num_columns = ReadBigEndian(image_file); + if (num_columns != MNISTImage::kSize) { + return tensorflow::errors::FailedPrecondition( + "Num Columns of Mnist Data File was ", num_columns, " expected 28"); + } + + // Iterate through the images, and create an MNISTImage struct for each + for (int i = 0; i < num_images; ++i) { + images.emplace_back(); + uint8_t img_buf[MNISTImage::kSize * MNISTImage::kSize]; + image_file.read(reinterpret_cast(&img_buf[0]), sizeof(img_buf)); + + // Convert the buffer into float MNISTImage + std::copy_n(img_buf, MNISTImage::kSize * MNISTImage::kSize, + images[i].buf.data()); + } + if (images.empty()) { + return tensorflow::errors::Internal("Error reading MNIST images"); + } + return tensorflow::Status::OK(); +} + +} // namespace mnist diff --git a/tftrt/examples/cpp/image-classification/mnist.h b/tftrt/examples/cpp/image-classification/mnist.h new file mode 100644 index 000000000..89a1e8385 --- /dev/null +++ b/tftrt/examples/cpp/image-classification/mnist.h @@ -0,0 +1,55 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef SAVED_MODEL_EXAMPLE_MNIST_H_ +#define SAVED_MODEL_EXAMPLE_MNIST_H_ + +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/lib/core/status.h" +#include +#include +#include + +namespace mnist { + +// MNISTImage contains the 28 x 28 pixel data for +// an MNIST picture, in row major order. +struct MNISTImage { + static constexpr int kSize = 28; + std::array buf; +}; + +// Prints an ASCII art representationof the MNISTImage, useful for debugging. +std::ostream &operator<<(std::ostream &, const MNISTImage &); + +// MNISTImageReader helps read a vector of MNISTImages +// from a file path to the MNIST Image dataset. Note +// that the file must already be decompressed. +class MNISTImageReader { +public: + MNISTImageReader(const std::string &mnist_path); + tensorflow::Status ReadMnistImages(); + // Converts an MNIST Image to a tensorflow Tensor. + tensorflow::Tensor MNISTImageToTensor(int offset, int batch_size); + + std::vector images; + +private: + std::string mnist_path_; +}; + +} // namespace mnist + +#endif // SAVED_MODEL_EXAMPLE_MNIST_H_ diff --git a/tftrt/examples/cpp/image-classification/mnist_train.py b/tftrt/examples/cpp/image-classification/mnist_train.py new file mode 100644 index 000000000..18d460951 --- /dev/null +++ b/tftrt/examples/cpp/image-classification/mnist_train.py @@ -0,0 +1,66 @@ +# Copyright 2021 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import tensorflow as tf +from tensorflow.python.compiler.tensorrt import trt_convert as trt +from tensorflow.python.saved_model import signature_constants +from tensorflow.python.saved_model import tag_constants +from tensorflow.python.framework import convert_to_constants +import numpy as np + +mnist = tf.keras.datasets.mnist + +(x_train, y_train), (x_test, y_test) = mnist.load_data() +x_train, x_test = x_train / 255.0, x_test / 255.0 + +model = tf.keras.models.Sequential([ + tf.keras.layers.Flatten(input_shape=(28, 28)), + tf.keras.layers.Dense(128, activation='relu'), + #tf.keras.layers.Dropout(0.2), + tf.keras.layers.Dense(10, activation='softmax') +]) + +model.compile(optimizer='adam', + loss='sparse_categorical_crossentropy', + metrics=['accuracy']) +model.summary() + +model.input_shape + +model.fit(x_train, y_train, epochs=5) +model.evaluate(x_test, y_test, verbose=2) + +model.save('mnist_model') + +print("Model trained and saved under mnist_model") + +# Load the model, save the variables and save it + +def get_func_from_saved_model(saved_model_dir): + saved_model_loaded = tf.saved_model.load( + saved_model_dir, tags=[tag_constants.SERVING]) + graph_func = saved_model_loaded.signatures[ + signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] + return graph_func, saved_model_loaded + +func, loaded_model = get_func_from_saved_model('mnist_model') + +# Create frozen func +frozen_func = convert_to_constants.convert_variables_to_constants_v2(func) +module = tf.Module() +module.myfunc = frozen_func +tf.saved_model.save(module,'mnist_model_frozen', signatures=frozen_func) + +print("Model frozen and saved under mnist_model_frozen")