TF模型tensorrt C++ API搭建部署
TF模型tensorrt C++ API搭建部署
PS
该文章主要参考了基于TensorRT C++ API 加速 TF 模型以及tensorrtx，记录face_landmark tf模型的tensorrt部署，具体代码参考https://github.com/ycdhqzhiai/face_landmarks_tensorrtx(仓库已经添加了pytorch实现)
1.ckpt转pb模型
git clone https://github.com/610265158/face_landmark
git checkout tf1
- 根据仓库里面转tflite的代码，简单修改下，可以直接生成pb模型
#-*-coding:utf-8-*-
import sys
sys.path.append('.')
import tensorflow as tf
import tensorflow.contrib.slim as slim
import math
from train_config import config as cfg
from lib.core.model.Mobilenet import mobilenet
from lib.helper import logger
###also u can change it to a specific model
free_graphmodel_folder = './model'
checkpoint = tf.train.get_checkpoint_state(model_folder).model_checkpoint_path
pretrained_model=checkpoint
saved_name ='./model/landmark_deploy.ckpt'
class trainner():
def __init__(self):
self.inputs=[]
self.outputs=[]
self.ite_num=1
def _wing_loss(self,landmarks, labels, w=10.0, epsilon=2.0):
"""
Arguments:
landmarks, labels: float tensors with shape [batch_size, landmarks]. landmarks means x1,x2,x3,x4...y1,y2,y3,y4 1-D
w, epsilon: a float numbers.
Returns:
a float tensor with shape [].
"""
with tf.name_scope('wing_loss'):
x = landmarks - labels
c = w * (1.0 - math.log(1.0 + w / epsilon))
absolute_x = tf.abs(x)
losses = tf.where(
tf.greater(w, absolute_x),
w * tf.log(1.0 + absolute_x / epsilon),
absolute_x - c
)
if cfg.TRAIN.ohem:
return losses
else:
loss = tf.reduce_mean(tf.reduce_mean(losses, axis=[1]), axis=0)
return loss
def tower_loss(self,scope, images, labels, L2_reg):
"""Calculate the total loss on a single tower running the model.
Args:
scope: unique prefix string identifying the CIFAR tower, e.g. 'tower_0'
images: Images. 4D tensor of shape [batch_size, height, width, 3].
labels: Labels. 1D tensor of shape [batch_size].
Returns:
Tensor of shape [] containing the total loss for a batch of data
"""
# Build the portion of the Graph calculating the losses. Note that we will
# assemble the total_loss using a custom function below.
#net_out = shufflenet_v2(images, L2_reg, False)
#net_out = resnet(images, L2_reg, False)
net_out = mobilenet(images, L2_reg, False)
loss, leye_loss, reye_loss, mouth_loss, leye_cla_accuracy,\
reye_cla_accuracy, mouth_cla_accuracy, l2_loss=calculate_loss(net_out,labels,scope)
return loss,leye_loss,reye_loss,mouth_loss,leye_cla_accuracy,reye_cla_accuracy,mouth_cla_accuracy, l2_loss
def build(self):
"""Train faces data for a number of epoch."""
with tf.Graph().as_default(), tf.device('/cpu:0'):
# Create a variable to count the number of train() calls. This equals the
# number of batches processed * FLAGS.num_gpus.
global_step = tf.get_variable(
'global_step', [],
initializer=tf.constant_initializer(0), dtype=tf.int32, trainable=False)
# Decay the learning rate
lr = tf.train.piecewise_constant(global_step,
cfg.TRAIN.lr_decay_every_step,
cfg.TRAIN.lr_value_every_step
)
keep_prob = tf.placeholder(tf.float32, name="keep_prob")
L2_reg = tf.placeholder(tf.float32, name="L2_reg")
images_place_holder_list = []
labels_place_holder_list = []
# Create an optimizer that performs gradient descent.
#opt = tf.train.AdamOptimizer(lr)
opt = tf.train.MomentumOptimizer(lr,momentum=0.9,use_nesterov=False)
# Get images and labels
weights_initializer = slim.xavier_initializer()
biases_initializer = tf.constant_initializer(0.)
biases_regularizer = tf.no_regularizer
weights_regularizer = tf.contrib.layers.l2_regularizer(L2_reg)
# Calculate the gradients for each model tower.
tower_grads = []
with tf.variable_scope(tf.get_variable_scope()):
for i in range(1):
with tf.device('/gpu:%d' % i):
with tf.name_scope('tower_%d' % (i)) as scope:
with slim.arg_scope([slim.model_variable, slim.variable], device='/cpu:0'):
images_ = tf.placeholder(tf.float32, [None, cfg.MODEL.hin, cfg.MODEL.win, 3], name="images")
labels_ = tf.placeholder(tf.float32, [None, cfg.MODEL.out_channel],name="labels")
images_place_holder_list.append(images_)
labels_place_holder_list.append(labels_)
with slim.arg_scope([slim.conv2d, slim.conv2d_in_plane, \
slim.conv2d_transpose, slim.separable_conv2d,
slim.fully_connected],
weights_regularizer=weights_regularizer,
biases_regularizer=biases_regularizer,
weights_initializer=weights_initializer,
biases_initializer=biases_initializer):
loss, leye_loss, reye_loss, mouth_loss,leye_cla_accuracy,reye_cla_accuracy,mouth_cla_accuracy, l2_loss = self.tower_loss(
scope, images_, labels_, L2_reg)
##use muti gpu ,large batch
if i == cfg.TRAIN.num_gpu - 1:
total_loss = tf.add_n([loss, leye_loss, reye_loss, mouth_loss, l2_loss])
else:
total_loss = tf.add_n([loss, leye_loss, reye_loss, mouth_loss])
# Reuse variables for the next tower.
tf.get_variable_scope().reuse_variables()
##when use batchnorm, updates operations only from the
## final tower. Ideally, we should grab the updates from all towers
# but these stats accumulate extremely fast so we can ignore the
# other stats from the other towers without significant detriment.
bn_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope=scope)
# Retain the summaries from the final tower.
#summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)
free_graph summaries = tf.get_collection('%smutiloss'%scope, scope)
# Calculate the gradients for the batch of data on this CIFAR tower.
grads = opt.compute_gradients(total_loss)
# Keep track of the gradients across all towers.
tower_grads.append(grads)
# We must calculate the mean of each gradient. Note that this is the
# synchronization point across all towers.
self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
# Build the summary operation from the last tower summaries.
# Build an initialization operation to run below.
init = tf.global_variables_initializer()
# Start running operations on the Graph. allow_soft_placement must be set to
# True to build towers on GPU, as some of the ops do not have GPU
# implementations.
tf_config = tf.ConfigProto(
allow_soft_placement=True,
log_device_placement=False)
tf_config.gpu_options.allow_growth = True
self.sess = tf.Session(config=tf_config)
self.sess.run(init)
#########################restore the params
variables_restore = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES)#,scope=cfg.MODEL.net_structure)
saver2 = tf.train.Saver(variables_restore)
saver2.restore(self.sess, pretrained_model)
logger.info('landmark_deploy saved')
self.saver.save(self.sess, save_path=saved_name)
self.sess.close()
train=trainner()
train.build()
将train_flage置为False，运行脚本重新生成ckpt模型
python net_work_for_tf_lite.py
- 去除auto_freeze脚本里面training参数直接转换
import sys
sys.path.append('.')
import os
import tensorflow as tf
from train_config import config as cfg
model_folder = cfg.MODEL.model_path
checkpoint = tf.train.get_checkpoint_state(model_folder)
##input_checkpoint
input_checkpoint = checkpoint.model_checkpoint_path
##input_graph
input_meta_graph = input_checkpoint + '.meta'
##output_node_names
output_node_names='tower_0/images,tower_0/prediction'
#output_graph
output_graph='./model/keypoints.pb'
print('excuted')
command="python tools/freeze.py --input_checkpoint %s --input_meta_graph %s --output_node_names %s --output_graph %s"\
%(input_checkpoint,input_meta_graph,output_node_names,output_graph)
os.system(command)
生成最终pb模型
python auto_freeze.py
2.将pb模型转换为wts文件
- 基于tf1
上面第一部分都是基于tf1将ckpt转换为pb模型，下面脚本是将pb转换为wts文件，参考了基于TensorRT C++ API 加速 TF 模型这篇文章，具体注意事项可以查看原文
import tensorflow as tf
from tensorflow.python.platform import gfile
import struct
import torch
from tensorflow.python.framework import tensor_util
# path to your .pb file
GRAPH_PB_PATH = './model/keypoints.pb'
GRAPH_WTS_PATH = './model/keypoints.wts'
with tf.Session() as sess:
print("load graph")
with gfile.FastGFile(GRAPH_PB_PATH, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
graph_nodes = [n for n in graph_def.node]
wts = [n for n in graph_nodes if n.op == 'Const']
dict = {}
for n in wts:
v = n.attr['value']
print(n.name)
ar = tensor_util.MakeNdarray(v.tensor)
dict[n.name] = torch.Tensor(ar)
f = open(GRAPH_WTS_PATH, 'w')
f.write("{}\n".format(len(dict.keys())))
for k, v in dict.items():
print('key: ', k)
print('value: ', v.shape)
if v.ndim == 4: # tf:NHWC trt:NCHW
v = v.transpose(3, 0).transpose(2, 1).transpose(3, 2)
vr = v.reshape(-1).cpu().numpy()
else:
vr = v.reshape(-1).cpu().numpy()
f.write("{} {}".format(k, len(vr)))
for vv in vr:
f.write(" ")
f.write(struct.pack(">f", float(vv)).hex())
f.write("\n")
- 基于tf2
源训练仓库tf2模型主要使用了tf.saved_model.save/tf.saved_model.load这2个api，用来保存和加载save_mode.pb模型，由于不是蛮会tf，查阅相关API，写了以下脚本转换wts文件
该脚本转换出来的wts文件，模型权重与netron对比是一致的，但是模型节点名称却很奇怪。没有接下来验证，感兴趣的可以验证下
3.根据wts 使用tensorRT API搭建网络
基于TensorRT C++ API 加速 TF 模型这篇文章中详细提到了各个层的使用到的API，我这里参考这篇文章，根据训练源码中shufflenetv2，重写了一遍，具体可以参考
https://github.com/ycdhqzhiai/face_landmarks_tensorrtx