最近看到一個(gè)巨牛的人工智能教程，分享一下給大家型雳。教程不僅是零基礎(chǔ)当凡，通俗易懂，而且非常風(fēng)趣幽默纠俭，像看小說一樣沿量！覺得太牛了，所以分享給大家冤荆。平時(shí)碎片時(shí)間可以當(dāng)小說看朴则，【點(diǎn)這里可以去膜拜一下大神的“小說”】。
在Tensorflow框架訓(xùn)練完成后匙赞，部署模型時(shí)希望對(duì)模型進(jìn)行壓縮佛掖。一種方案是前面文字介紹的方法《【Ubuntu】Tensorflow對(duì)訓(xùn)練后的模型做8位（uint8）量化轉(zhuǎn)換》妖碉。另一種方法是半浮點(diǎn)量化，今天我們主要介紹如何通過修改Tensorflow的pb文件中的計(jì)算節(jié)點(diǎn)和常量（const）芥被，將float32數(shù)據(jù)類型的模型大小壓縮減半為float16數(shù)據(jù)類型的模型欧宜。
1 加載pb模型
封裝函數(shù)，加載pb模型：
def load_graph(model_path):
graph = tf.Graph()
with graph.as_default():
graph_def = tf.GraphDef()
if model_path.endswith("pb"):
with open(model_path, "rb") as f:
graph_def.ParseFromString(f.read())
else:
with open(model_path, "r") as pf:
text_format.Parse(pf.read(), graph_def)
tf.import_graph_def(graph_def, name="")
sess = tf.Session(graph=graph)
ops=graph.get_operations()
for op in ops:
print(op.name)
return sess
2 重寫B(tài)atchNorm
由于BatchNorm對(duì)精度比較敏感拴魄，需要保持float32類型冗茸，因此BatchNorm需要特殊處理。
#用FusedBatchNormV2替換FusedBatchNorm匹中，以保證反向梯度下降計(jì)算時(shí)使用的是float
def rewrite_batch_norm_node_v2(node, graph_def, target_type='fp16'):
if target_type == 'fp16':
dtype = types_pb2.DT_HALF
elif target_type == 'fp64':
dtype = types_pb2.DT_DOUBLE
else:
dtype = types_pb2.DT_FLOAT
new_node = graph_def.node.add()
new_node.op = "FusedBatchNormV2"
new_node.name = node.name
new_node.input.extend(node.input)
new_node.attr["U"].CopyFrom(attr_value_pb2.AttrValue(type=types_pb2.DT_FLOAT))
for attr in list(node.attr.keys()):
if attr == "T":
node.attr[attr].type = dtype
new_node.attr[attr].CopyFrom(node.attr[attr])
print("rewrite fused_batch_norm done!")
3 Graph轉(zhuǎn)換
重新構(gòu)造graph夏漱，參數(shù)從原始pb的graph中拷貝，并轉(zhuǎn)為float16
def convert_graph_to_fp16(model_path, save_path, name, as_text=False, target_type='fp16', input_name=None, output_names=None):
#生成新的圖數(shù)據(jù)類型
if target_type == 'fp16':
dtype = types_pb2.DT_HALF
elif target_type == 'fp64':
dtype = types_pb2.DT_DOUBLE
else:
dtype = types_pb2.DT_FLOAT
#加載需要轉(zhuǎn)換的模型
source_sess = load_graph(model_path)
source_graph_def = source_sess.graph.as_graph_def()
#創(chuàng)建新的模圖對(duì)象
target_graph_def = graph_pb2.GraphDef()
target_graph_def.versions.CopyFrom(source_graph_def.versions)
#對(duì)加載的模型遍歷計(jì)算節(jié)點(diǎn)
for node in source_graph_def.node:
# 對(duì)FusedBatchNorm計(jì)算節(jié)點(diǎn)替換為FusedBatchNormV2
if node.op == "FusedBatchNorm":
rewrite_batch_norm_node_v2(node, target_graph_def, target_type=target_type)
continue
# 復(fù)制計(jì)算節(jié)點(diǎn)
new_node = target_graph_def.node.add()
new_node.op = node.op
new_node.name = node.name
new_node.input.extend(node.input)
#對(duì)attrs屬性進(jìn)行復(fù)制顶捷，attrs屬性主要關(guān)注
attrs = list(node.attr.keys())
# BatchNorm屬性保持不變
if ("BatchNorm" in node.name) or ('batch_normalization' in node.name):
for attr in attrs:
new_node.attr[attr].CopyFrom(node.attr[attr])
continue
# 除了BatchNorm以外其他計(jì)算節(jié)點(diǎn)的屬性單獨(dú)
for attr in attrs:
# 對(duì)指定的計(jì)算節(jié)點(diǎn)保持不變
if node.name in keep_fp32_node_name:
new_node.attr[attr].CopyFrom(node.attr[attr])
continue
#將Float類型修改為設(shè)置的目標(biāo)類型
if node.attr[attr].type == types_pb2.DT_FLOAT:
# modify node dtype
node.attr[attr].type = dtype
#重點(diǎn)關(guān)注value挂绰，weights都是保存在value屬性中
if attr == "value":
tensor = node.attr[attr].tensor
if tensor.dtype == types_pb2.DT_FLOAT:
# if float_val exists
if tensor.float_val:
float_val = tf.make_ndarray(node.attr[attr].tensor)
new_node.attr[attr].tensor.CopyFrom(tf.make_tensor_proto(float_val, dtype=dtype))
continue
# if tensor content exists
if tensor.tensor_content:
tensor_shape = [x.size for x in tensor.tensor_shape.dim]
tensor_weights = tf.make_ndarray(tensor)
# reshape tensor
tensor_weights = np.reshape(tensor_weights, tensor_shape)
tensor_proto = tf.make_tensor_proto(tensor_weights, dtype=dtype)
new_node.attr[attr].tensor.CopyFrom(tensor_proto)
continue
new_node.attr[attr].CopyFrom(node.attr[attr])
# transform graph
if output_names:
if not input_name:
input_name = []
transforms = ["strip_unused_nodes"]
target_graph_def = TransformGraph(target_graph_def, input_name, output_names, transforms)
# write graph_def to model
tf.io.write_graph(target_graph_def, logdir=save_path, name=name, as_text=as_text)
print("Converting done ...")
4 完整的代碼
import tensorflow as tf
from tensorflow.core.framework import types_pb2, graph_pb2, attr_value_pb2
from tensorflow.tools.graph_transforms import TransformGraph
from google.protobuf import text_format
import numpy as np
# object detection api input and output nodes
input_name = "input_tf"
output_names = ["output:0"]
keep_fp32_node_name = []
def load_graph(model_path):
graph = tf.Graph()
with graph.as_default():
graph_def = tf.GraphDef()
if model_path.endswith("pb"):
with open(model_path, "rb") as f:
graph_def.ParseFromString(f.read())
else:
with open(model_path, "r") as pf:
text_format.Parse(pf.read(), graph_def)
tf.import_graph_def(graph_def, name="")
sess = tf.Session(graph=graph)
ops=graph.get_operations()
for op in ops:
print(op.name)
return sess
#用FusedBatchNormV2替換FusedBatchNorm，以保證反向梯度下降計(jì)算時(shí)使用的是float
def rewrite_batch_norm_node_v2(node, graph_def, target_type='fp16'):
if target_type == 'fp16':
dtype = types_pb2.DT_HALF
elif target_type == 'fp64':
dtype = types_pb2.DT_DOUBLE
else:
dtype = types_pb2.DT_FLOAT
new_node = graph_def.node.add()
new_node.op = "FusedBatchNormV2"
new_node.name = node.name
new_node.input.extend(node.input)
new_node.attr["U"].CopyFrom(attr_value_pb2.AttrValue(type=types_pb2.DT_FLOAT))
for attr in list(node.attr.keys()):
if attr == "T":
node.attr[attr].type = dtype
new_node.attr[attr].CopyFrom(node.attr[attr])
print("rewrite fused_batch_norm done!")
def convert_graph_to_fp16(model_path, save_path, name, as_text=False, target_type='fp16', input_name=None, output_names=None):
#生成新的圖數(shù)據(jù)類型
if target_type == 'fp16':
dtype = types_pb2.DT_HALF
elif target_type == 'fp64':
dtype = types_pb2.DT_DOUBLE
else:
dtype = types_pb2.DT_FLOAT
#加載需要轉(zhuǎn)換的模型
source_sess = load_graph(model_path)
source_graph_def = source_sess.graph.as_graph_def()
#創(chuàng)建新的模圖對(duì)象
target_graph_def = graph_pb2.GraphDef()
target_graph_def.versions.CopyFrom(source_graph_def.versions)
#對(duì)加載的模型遍歷計(jì)算節(jié)點(diǎn)
for node in source_graph_def.node:
# 對(duì)FusedBatchNorm計(jì)算節(jié)點(diǎn)替換為FusedBatchNormV2
if node.op == "FusedBatchNorm":
rewrite_batch_norm_node_v2(node, target_graph_def, target_type=target_type)
continue
# 復(fù)制計(jì)算節(jié)點(diǎn)
new_node = target_graph_def.node.add()
new_node.op = node.op
new_node.name = node.name
new_node.input.extend(node.input)
#對(duì)attrs屬性進(jìn)行復(fù)制服赎，attrs屬性主要關(guān)注
attrs = list(node.attr.keys())
# BatchNorm屬性保持不變
if ("BatchNorm" in node.name) or ('batch_normalization' in node.name):
for attr in attrs:
new_node.attr[attr].CopyFrom(node.attr[attr])
continue
# 除了BatchNorm以外其他計(jì)算節(jié)點(diǎn)的屬性單獨(dú)
for attr in attrs:
# 對(duì)指定的計(jì)算節(jié)點(diǎn)保持不變
if node.name in keep_fp32_node_name:
new_node.attr[attr].CopyFrom(node.attr[attr])
continue
#將Float類型修改為設(shè)置的目標(biāo)類型
if node.attr[attr].type == types_pb2.DT_FLOAT:
# modify node dtype
node.attr[attr].type = dtype
#重點(diǎn)關(guān)注value葵蒂，weights都是保存在value屬性中
if attr == "value":
tensor = node.attr[attr].tensor
if tensor.dtype == types_pb2.DT_FLOAT:
# if float_val exists
if tensor.float_val:
float_val = tf.make_ndarray(node.attr[attr].tensor)
new_node.attr[attr].tensor.CopyFrom(tf.make_tensor_proto(float_val, dtype=dtype))
continue
# if tensor content exists
if tensor.tensor_content:
tensor_shape = [x.size for x in tensor.tensor_shape.dim]
tensor_weights = tf.make_ndarray(tensor)
# reshape tensor
tensor_weights = np.reshape(tensor_weights, tensor_shape)
tensor_proto = tf.make_tensor_proto(tensor_weights, dtype=dtype)
new_node.attr[attr].tensor.CopyFrom(tensor_proto)
continue
new_node.attr[attr].CopyFrom(node.attr[attr])
# transform graph
if output_names:
if not input_name:
input_name = []
transforms = ["strip_unused_nodes"]
target_graph_def = TransformGraph(target_graph_def, input_name, output_names, transforms)
# write graph_def to model
tf.io.write_graph(target_graph_def, logdir=save_path, name=name, as_text=as_text)
print("Converting done ...")
save_path = "test"
name = "output_fp16.pb"
model_path="test.pb"
as_text = False
target_type = 'fp16'
convert_graph_to_fp16(model_path, save_path, name, as_text=as_text, target_type=target_type, input_name=input_name, output_names=output_names)
# 測(cè)試一下轉(zhuǎn)換后的模型是否能夠加載
sess = load_graph(save_path+"/"+name)
