tensorflow实现resnet50（训练+测试+模型转换）
本章使用tensorflow训练resnet50，使用手写数字图片作为数据集。
数据集：
代码工程：
1.train.py
import argparse
import cv2
import tensorflow as tf
# from create_model import resnet_v2_50
from create_model import resnet_v2_50
import numpy as np
from data_loader import get_data, get_data_list
from sklearn.metrics import accuracy_score
def txt_save(data, output_file):
file = open(output_file, 'a')
for i in data:
s = str(i) + '\n'
file.write(s)
file.close()
def get_parms():
parser = argparse.ArgumentParser(description='')
parser.add_argument('--train_data', type=str, default="dataset/train_data.txt")
parser.add_argument('--test_data', type=str, default='data/test/')
parser.add_argument('--checkpoint_dir', type=str, default='./model/')
parser.add_argument('--epoch', type=int, default=5)
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--save_epoch', type=int, default=1)
args=parser.parse_args()
return args
args = get_parms()
inputs=tf.placeholder(tf.float32,(None,28,28,1), name='input_images')
labels = tf.placeholder(tf.int64, [None, 10])
net,endpoins=resnet_v2_50(inputs,10) #['predictions']
# with tf.variable_scope('finetune'):
logit = tf.nn.softmax(net)[0][0]
pred=tf.argmax(logit,1)
correct_predicion=tf.equal(tf.argmax(logit,1),tf.argmax(labels,1))
accuracy=tf.reduce_mean(tf.cast(correct_predicion,'float'))
cross_entropy = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logit)
cross_entropy_cost = tf.reduce_mean(cross_entropy)
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy_cost) #Resnet_v2_50和自己构建的模型层都训练
# 开始训练
saver=tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
img_list, label_list=get_data_list(args.train_data)
for i in range(args.epoch):
loss_list=[]
acc_list=[]
for j in range(int(len(label_list)/args.batch_size)):
# for j in range(10):
data, true_label=get_data(img_list, label_list, args.batch_size)
_, loss, acc=sess.run([train_step, cross_entropy_cost, accuracy], feed_dict={inputs:data, labels:true_label})
# _=sess.run([train_step], feed_dict={inputs:data, labels:true_label})
# print(loss, acc)
# a,b=sess.run([pred,logit], feed_dict={inputs:data, labels:true_label})
# print(a,b)
print('epoch:',i, 'loss:',np.mean(loss), "acc:", acc)
if i % args.save_epoch==0:
# saver.save(sess,"model/model.ckpt",global_step=i)
saver.save(sess, "model/model")
# tensorboard --logdir=d:/log --host=127.0.0.1
init = tf.initialize_all_variables()
merged = tf.summary.merge_all()
writer = tf.summary.FileWriter("d:/log/",sess.graph) #目录结构尽量简单，复杂了容易出现找不到文件，原因不清楚
sess.run(init)
# variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
# for i in variables:
# print(i) # 打印
# txt_save(variables, "model/log.txt") # 保存txt 二选一
# import tensorflow as tf
# img_list, label_list=get_data_list(args.train_data)
# data, true_label=get_data(img_list, label_list, args.batch_size)
# with tf.Session() as sess:
# saver = tf.train.import_meta_graph('model/model.meta')
# saver.restore(sess,tf.train.latest_checkpoint('model/'))
# pred,logit=sess.run([pred,logit], feed_dict={inputs:data, labels:true_label})
# print(pred)
# ##Model has been restored. Above statement will print the saved value
2.create_model.py
import tensorflow as tf
import os
import numpy as np
import cv2
import argparse
def model(inputs):
w1=tf.Variable(tf.random_normal([3, 3, 1, 32], stddev=0.01))
w2=tf.Variable(tf.random_normal([3, 3, 32, 64], stddev=0.01))
w3=tf.Variable(tf.random_normal([3, 3, 64, 128], stddev=0.01))
w4=tf.Variable(tf.random_normal([2048, 625], stddev=0.01))
w5=tf.Variable(tf.random_normal([625, 10], stddev=0.01))
l1_conv=tf.nn.conv2d(inputs, w1, strides=[1, 1, 1, 1], padding='SAME')
l1_relu=tf.nn.relu(l1_conv)
l1_pool=tf.nn.max_pool(l1_relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
l1_drop = tf.nn.dropout(l1_pool, 0.5)
l2_conv=tf.nn.conv2d(l1_drop, w2, strides=[1, 1, 1, 1], padding='SAME')
l2_relu=tf.nn.relu(l2_conv)
l2_pool=tf.nn.max_pool(l2_relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
l2_drop = tf.nn.dropout(l2_pool, 0.5)
l3_conv=tf.nn.conv2d(l2_drop, w3, strides=[1, 1, 1, 1], padding='SAME')
l3_relu=tf.nn.relu(l3_conv)
l3_pool=tf.nn.max_pool(l3_relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
l3_out = tf.reshape(l3_pool, [-1, 2048])
l3_drop = tf.nn.dropout(l3_out, 0.5)
l4 = tf.nn.relu(tf.matmul(l3_drop, w4))
l4 = tf.nn.dropout(l4, 0.5)
out = tf.matmul(l4, w5)
return out
def model2(inputs):
w1=tf.Variable(tf.random_normal([5, 5, 1, 6], stddev=0.01))
b1 = tf.Variable(tf.truncated_normal([6]))
w2=tf.Variable(tf.random_normal([5, 5, 6, 16], stddev=0.01))
b2 = tf.Variable(tf.truncated_normal([16]))
w3=tf.Variable(tf.random_normal([5, 5, 16, 120], stddev=0.01))
b3 = tf.Variable(tf.truncated_normal([120]))
w4 = tf.Variable(tf.truncated_normal([7 * 7 * 120, 80]))
b4 = tf.Variable(tf.truncated_normal([80]))
w5 = tf.Variable(tf.truncated_normal([80, 10]))
b5 = tf.Variable(tf.truncated_normal([10]))
l1_conv=tf.nn.conv2d(inputs, w1, strides=[1, 1, 1, 1], padding='SAME')
l1_sigmoid=tf.nn.sigmoid(l1_conv+b1)
l1_pool=tf.nn.max_pool(l1_sigmoid, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
l2_conv=tf.nn.conv2d(l1_pool, w2, strides=[1, 1, 1, 1], padding='SAME')
l2_sigmoid=tf.nn.sigmoid(l2_conv+b2)
l2_pool=tf.nn.max_pool(l2_sigmoid, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
l3_conv=tf.nn.conv2d(l2_pool, w3, strides=[1, 1, 1, 1], padding='SAME')
l3_sigmoid=tf.nn.sigmoid(l3_conv+b3)
l3_out = tf.reshape(l3_sigmoid, [-1, 7*7*120])
l4 = tf.nn.sigmoid(tf.matmul(l3_out, w4)+b4)
out = tf.nn.softmax(tf.matmul(l4, w5) + b5)
return out
from datetime import datetime
import time
import math
import collections
import tensorflow as tf
slim = tf.contrib.slim
# 使用collections.namedtuple设计ResNet的Block模块
# scope参数是block的名称
# unit_fn是功能单元（如残差单元）
# args是一个列表，如([256, 64, 1]) X 2 + [256, 64, 2]),代表两个（256, 64, 1）单元
# 和一个（256, 64, 2）单元
Block = collections.namedtuple("Block", ['scope', 'unit_fn', 'args'])
# 定义下采样的方法，通过max_pool2d实现
def subsample(inputs, factor, scope=None):
if factor == 1:
return inputs
else:
return slim.max_pool2d(inputs, [1, 1], stride=factor, scope=scope)
# 定义一个创建卷积层的函数
def conv2d_same(inputs, num_outputs, kernel_size, stride, scope=None):
if stride == 1:
return slim.conv2d(inputs, num_outputs, kernel_size, stride=1,
padding='SAME', scope=scope)
else:
pad_total = kernel_size - 1
pad_beg = pad_total // 2
pad_end = pad_total - pad_beg
inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end],
[pad_beg, pad_end], [0, 0]])
return slim.conv2d(inputs, num_outputs, kernel_size, stride=stride,
padding='VALID', scope=scope)
# 定义堆叠的block函数
@slim.add_arg_scope
def stack_blocks_dense(net, blocks, outputs_collections=None):
for block in blocks:
with tf.variable_scope(block.scope, 'block', [net]) as sc:
for i, unit in enumerate(block.args):
with tf.variable_scope('unit_%d' % (i + 1), values=[net]):
unit_depth, unit_depth_bottleneck, unit_stride = unit
net = block.unit_fn(net,
depth=unit_depth,
depth_bottleneck=unit_depth_bottleneck,
stride=unit_stride)
net = slim.utils.collect_named_outputs(outputs_collections, sc.name,net)
return net
# 用于设定默认值
def resnet_arg_scope(is_training=True,
weight_decay=0.0001,
batch_norm_decay=0.997,
batch_norm_epsilon=1e-5,
batch_norm_scale=True):
batch_norm_params = {
'is_training': is_training,
'decay': batch_norm_decay,
'epsilon': batch_norm_epsilon,
'scale': batch_norm_scale,
'updates_collections': tf.GraphKeys.UPDATE_OPS,
}
with slim.arg_scope(
[slim.conv2d],
weights_regularizer=slim.l2_regularizer(weight_decay),
weights_initializer=slim.variance_scaling_initializer(),
activation_fn=tf.nn.relu,
normalizer_fn=slim.batch_norm,
normalizer_params=batch_norm_params
):
with slim.arg_scope([slim.batch_norm], **batch_norm_params):
with slim.arg_scope([slim.max_pool2d], padding='SAME') as arg_sc:
return arg_sc
# 定义残差学习单元
@slim.add_arg_scope
def bottleneck(inputs, depth, depth_bottleneck, stride,
outputs_collections=None, scope=None):
with tf.variable_scope(scope, 'bottleneck_v2', [inputs]) as sc:
depth_in = slim.utils.last_dimension(inputs.get_shape(), min_rank=4)
preact = slim.batch_norm(inputs, activation_fn=tf.nn.relu,
scope='preact')
# shortcut为直连的X
if depth == depth_in:
shortcut = subsample(inputs, stride, 'shortcut')
else:
shortcut = slim.conv2d(preact, depth, [1, 1], stride=stride,
normalizer_fn=None, activation_fn=None,
scope='shortcut')
residual = slim.conv2d(preact, depth_bottleneck, [1, 1], stride=1,
scope='conv1')
residual = conv2d_same(residual, depth_bottleneck, 3, stride,
scope='conv2')
residual = slim.conv2d(residual, depth, [1, 1], stride=1,
normalizer_fn=None, activation_fn=None,
scope='conv3')
# 将直连的X加到残差上，得到output
output = shortcut + residual
return slim.utils.collect_named_outputs(outputs_collections,
sc.name, output)
# 定义ResNet的主函数
def resnet_v2(inputs,
blocks,
num_classes=None,
global_pool=True,
include_root_block=True,
reuse=None,
scope=None):
with tf.variable_scope(scope, 'resnet_v2', [inputs], reuse=reuse) as sc:
end_points_collection = sc.original_name_scope + '_end_points'
with slim.arg_scope([slim.conv2d, bottleneck,stack_blocks_dense],outputs_collections=end_points_collection):
net = inputs
if include_root_block:
with slim.arg_scope([slim.conv2d], activation_fn=None,normalizer_fn=None):
net = conv2d_same(net, 64, 7, stride=2, scope='conv1')
net = slim.max_pool2d(net, [3, 3], stride=2, scope='pool1')
net = stack_blocks_dense(net, blocks)
net = slim.batch_norm(net, activation_fn=tf.nn.relu, scope='postnorm')
if global_pool:
net = tf.reduce_mean(net, [1, 2], name='pool5', keep_dims=True)
if num_classes is not None:
net = slim.conv2d(net, num_classes, [1, 1], activation_fn=None,normalizer_fn=None, scope='logits')
end_points = slim.utils.convert_collection_to_dict(end_points_collection)
if num_classes is not None:
end_points['predictions'] = slim.softmax(net, scope='predictions')
return net, end_points
# 定义50层的ResNet
def resnet_v2_50(inputs,
num_classes=None,
global_pool=True,
reuse=None,
scope='resnet_v2_50'):
blocks = [
Block('block1', bottleneck, [(256, 64, 1)] * 2 + [(256, 64, 2)]),
Block('block2', bottleneck, [(512, 128, 1)] * 3 + [(512, 128, 2)]),
Block('block3', bottleneck, [(1024, 256, 1)] * 5 + [(1024, 256, 2)]),
Block('block4', bottleneck, [(2048, 512, 1)] * 3)]
return resnet_v2(inputs, blocks, num_classes, global_pool,include_root_block=True, reuse=reuse, scope=scope)3.data_loader.py
import cv2
import os
import numpy as np
import random
# f1=open('dataset/train_data.txt','w+')
# path='dataset/'
# for file in os.listdir(path):
# if file.endswith('png'):
# line=path+file+' 1'+'\n'
# f1.write(line)
# print(line)
def one_hot(data, num_classes):
return np.squeeze(np.eye(num_classes)[data.reshape(-1)])
def get_data_list(path):
f1=open(path,'r')
lines=f1.readlines()
img_list=[]
label_list=[]
for line in lines:
label=int(line.strip().split(" ")[1])
label=one_hot(np.array(label),10)
label_list.append(label)
file_name=line.strip().split(" ")[0]
img=cv2.imread(file_name, 0)
img=np.reshape(img,[28,28,1])
# print(img.shape)
img_list.append(img)
return img_list, label_list
def get_data(img_list, label_list, batch_size):
lens=len(label_list)
random_nums=random.sample(range(lens),lens)
nums=random_nums[0:batch_size]
# print(nums)
data=[]
label=[]
for index in nums:
data.append(img_list[index])
label.append(label_list[index])
return np.array(data), np.array(label)
# batch_size=1
# path="dataset/mnist/train/train_data.txt"
# img_list, label_list=get_data_list(path)
# print(len(img_list), len(label_list))
# data, label=get_data(img_list, label_list, batch_size)
# print(type(data[0]),label)4.查看ckpt网络.py
from tensorflow.python import pywrap_tensorflow
import os
import tensorflow as tf
from tensorflow.python.platform import gfile
# checkpoint_path = os.path.join('model/model.ckpt')
# reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
# var_to_shape_map = reader.get_variable_to_shape_map()
# for key in var_to_shape_map:
# print('tensor_name: ', key)
ckpt_path = os.path.join('model/model')
saver = tf.train.import_meta_graph(ckpt_path+'.meta',clear_devices=True)
graph = tf.get_default_graph()
with tf.Session( graph=graph) as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess,ckpt_path)
# tensorboard --logdir=d:/log --host=127.0.0.1
init = tf.initialize_all_variables()
merged = tf.summary.merge_all()
writer = tf.summary.FileWriter("d:/log/",sess.graph) #目录结构尽量简单，复杂了容易出现找不到文件，原因不清楚
sess.run(init)5.ckpt2pb.py
import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.python import pywrap_tensorflow
def freeze_graph(cpkt_path, pb_path):
checkpoint = tf.train.get_checkpoint_state("model") #检查目录下ckpt文件状态是否可用
cpkt_path2 = checkpoint.model_checkpoint_path #得ckpt文件路径
print("gg:",checkpoint,cpkt_path2)
# 指定输出的节点名称,该节点名称必须是原模型中存在的节点
output_node_names = "resnet_v2_50/logits/BiasAdd"
saver = tf.train.import_meta_graph(cpkt_path + '.meta', clear_devices=True)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
# feature_data_list = input_graph_def.get_operation_by_name('resnet_v2_50/conv1').outputs[0]
# input_image=tf.placeholder(None,28,28,1)
with tf.Session() as sess:
saver.restore(sess, cpkt_path) # 恢复图并得到数据
pb_path_def = graph_util.convert_variables_to_constants( # 模型持久化，将变量值固定
sess=sess,
input_graph_def=input_graph_def, # 等于:sess.graph_def
output_node_names=output_node_names.split(",")) # 如果有多个输出节点，以逗号隔开
# print(pb_path_def)
with tf.gfile.GFile(pb_path, 'wb') as fgraph:
fgraph.write(pb_path_def.SerializeToString())
# with tf.io.gfile.GFile(pb_path, "wb") as f: # 保存模型
# f.write(pb_path_def.SerializeToString()) # 序列化输出
print("%d ops in the final graph." % len(pb_path_def.node)) # 得到当前图有几个操作节点
if __name__ == '__main__':
# 输入路径(cpkt)
cpkt_path = 'model/model'
# 输出路径(pb模型)pb_path_def
pb_path = "model/test.pb"
# 模型转换
freeze_graph(cpkt_path, pb_path)
# # 查看节点名称：
# reader = pywrap_tensorflow.NewCheckpointReader(cpkt_path)
# var_to_shape_map = reader.get_variable_to_shape_map()
# for key in var_to_shape_map:
# print("tensor_name: ", key)
# # 查看某个指定节点的权重
# reader = pywrap_tensorflow.NewCheckpointReader(cpkt_path)
# var_to_shape_map = reader.get_variable_to_shape_map()
# w0 = reader.get_tensor("finetune/dense_1/bias")
# print(w0.shape, type(w0))
# print(w0[0])
# with tf.Session() as sess:
# # 加载模型定义的graph
# saver = tf.train.import_meta_graph('model/model.meta')
# # 方式一:加载指定文件夹下最近保存的一个模型的数据
# saver.restore(sess, tf.train.latest_checkpoint('model/'))
# # 方式二:指定具体某个数据，需要注意的是，指定的文件不要包含后缀
# # saver.restore(sess, os.path.join(path, 'model.ckpt-1000'))
# # 查看模型中的trainable variables
# tvs = [v for v in tf.trainable_variables()]
# for v in tvs:
# print(v.name)
# # print(sess.run(v))
# # # 查看模型中的所有tensor或者operations
# # gv = [v for v in tf.global_variables()]
# # for v in gv:
# # print(v.name)
# # # 获得几乎所有的operations相关的tensor
# # ops = [o for o in sess.graph.get_operations()]
# # for o in ops:
# # print(o.name)6.pb2pbtxt.py (pb和pbtxt互转，修改input_shape)
修改pbtxt文件，把动态shape改成静态的，再转回pb模型
import tensorflow as tf
from tensorflow.python.platform import gfile
from google.protobuf import text_format
def convert_pb_to_pbtxt(root_path, pb_path, pbtxt_path):
with gfile.FastGFile(root_path+pb_path, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
tf.train.write_graph(graph_def, root_path, pbtxt_path, as_text=True)
return
def convert_pbtxt_to_pb(root_path, pb_path, pbtxt_path):
with tf.gfile.FastGFile(root_path+pbtxt_path, 'r') as f:
graph_def = tf.GraphDef()
file_content = f.read()
# Merges the human-readable string in `file_content` into `graph_def`.
text_format.Merge(file_content, graph_def)
tf.train.write_graph(graph_def, root_path, pb_path, as_text=False)
return
if __name__ == '__main__':
# 模型路径
root_path = "model/"
pb_path = "test.pb"
pbtxt_path = "test.pbtxt"
# 模型转换
convert_pb_to_pbtxt(root_path, pb_path, pbtxt_path)
# convert_pbtxt_to_pb(root_path, pb_path, pbtxt_path)7.test_pb.py
import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.python import pywrap_tensorflow
import numpy as np
def recognize(jpg_path, pb_file_path):
with tf.Graph().as_default():
output_graph_def = tf.GraphDef()
with open(pb_file_path, "rb") as f:#主要步骤即为以下标出的几步,1、2步即为读取图
output_graph_def.ParseFromString(f.read())# 1.将模型文件解析为二进制放进graph_def对象
_ = tf.import_graph_def(output_graph_def, name="")# 2.import到当前图
with tf.Session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
graph = tf.get_default_graph()# 3.获得当前图
# # 4.get_tensor_by_name获取需要的节点
# x = graph.get_tensor_by_name("IteratorGetNext_1:0")
# y_out = graph.get_tensor_by_name("resnet_v1_50_1/predictions/Softmax:0")
x = graph.get_tensor_by_name("input_images:0")
y_out = graph.get_tensor_by_name("resnet_v2_50/logits/BiasAdd:0")
img=np.random.normal(size=(1, 28, 28, 1))
# img=cv2.imread(jpg_path)
# img=cv2.resize(img, (224, 224))
# img=np.reshape(img,(1,224,224,3))
# print(img.shape)
#执行
output = sess.run(y_out, feed_dict={x:img})
pred=np.argmax(output[0][0], axis=1)
print("预测结果：", output.shape, output, "预测label:", pred)
# prediction_labels = np.argmax(test_y_out, axis=2)
# print(prediction_labels.shape, prediction_labels)
recognize("dataset/00000.PNG", "model/test.pb")