devbox
我家自动化
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

tensorflow实现resnet50（训练+测试+模型转换）_tensorflow resnet50

作者：我家自动化 | 2024-02-16 11:35:58

赞

踩

tensorflow resnet50

本章使用tensorflow训练resnet50，使用手写数字图片作为数据集。

数据集：

代码工程：

1.train.py

  1. import argparse
  2. import cv2
  3. import tensorflow as tf
  4. # from create_model import resnet_v2_50
  5. from create_model import resnet_v2_50
  6. import numpy as np
  7. from data_loader import get_data, get_data_list
  8. from sklearn.metrics import accuracy_score
  9.  
  10. def txt_save(data, output_file):
  11.     file = open(output_file, 'a')
  12.     for i in data:
  13.         s = str(i) + '\n'
  14.         file.write(s)
  15.     file.close()
  16.  
  17. def get_parms():
  18.     parser = argparse.ArgumentParser(description='')
  19.  
  20.     parser.add_argument('--train_data', type=str, default="dataset/train_data.txt")
  21.     parser.add_argument('--test_data', type=str, default='data/test/')
  22.  
  23.     parser.add_argument('--checkpoint_dir', type=str, default='./model/')
  24.     parser.add_argument('--epoch', type=int, default=5)
  25.     parser.add_argument('--batch_size', type=int, default=1)
  26.     parser.add_argument('--save_epoch', type=int, default=1)
  27.     
  28.     args=parser.parse_args()
  29.  
  30.     return args
  31.  
  32.  
  33. args = get_parms()
  34. inputs=tf.placeholder(tf.float32,(None,28,28,1), name='input_images')
  35. labels = tf.placeholder(tf.int64, [None, 10])
  36. net,endpoins=resnet_v2_50(inputs,10)  #['predictions']
  37.  
  38.  
  39. # with tf.variable_scope('finetune'):
  40. logit = tf.nn.softmax(net)[0][0]
  41. pred=tf.argmax(logit,1)
  42.  
  43. correct_predicion=tf.equal(tf.argmax(logit,1),tf.argmax(labels,1))
  44. accuracy=tf.reduce_mean(tf.cast(correct_predicion,'float'))
  45.  
  46. cross_entropy = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logit)
  47. cross_entropy_cost = tf.reduce_mean(cross_entropy)
  48.  
  49. train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy_cost) #Resnet_v2_50和自己构建的模型层都训练
  50.  
  51.  
  52. # 开始训练
  53. saver=tf.train.Saver()
  54. with tf.Session() as sess:
  55.     sess.run(tf.global_variables_initializer())
  56.  
  57.     img_list, label_list=get_data_list(args.train_data)
  58.     for i in range(args.epoch):
  59.         loss_list=[]
  60.         acc_list=[]
  61.         for j in range(int(len(label_list)/args.batch_size)):
  62.         # for j in range(10):
  63.             data, true_label=get_data(img_list, label_list, args.batch_size)
  64.             _, loss, acc=sess.run([train_step, cross_entropy_cost, accuracy], feed_dict={inputs:data, labels:true_label})
  65.             # _=sess.run([train_step], feed_dict={inputs:data, labels:true_label})
  66.             # print(loss, acc)
  67.  
  68.             # a,b=sess.run([pred,logit], feed_dict={inputs:data, labels:true_label})
  69.             # print(a,b)
  70.         
  71.         print('epoch:',i, 'loss:',np.mean(loss), "acc:", acc)
  72.         if i % args.save_epoch==0:
  73.             # saver.save(sess,"model/model.ckpt",global_step=i)
  74.             saver.save(sess, "model/model")
  75.  
  76.  
  77.     # tensorboard --logdir=d:/log  --host=127.0.0.1
  78.     init = tf.initialize_all_variables()
  79.     merged = tf.summary.merge_all()
  80.     writer = tf.summary.FileWriter("d:/log/",sess.graph) #目录结构尽量简单，复杂了容易出现找不到文件，原因不清楚
  81.     sess.run(init)
  82.  
  83.     # variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
  84.     # for i in variables:
  85.     #    print(i)     # 打印
  86.     #    txt_save(variables, "model/log.txt")  # 保存txt   二选一
  87.  
  88.  
  89. # import tensorflow as tf
  90.  
  91. # img_list, label_list=get_data_list(args.train_data)
  92. # data, true_label=get_data(img_list, label_list, args.batch_size)
  93. # with tf.Session() as sess:
  94. #     saver = tf.train.import_meta_graph('model/model.meta')
  95. #     saver.restore(sess,tf.train.latest_checkpoint('model/'))
  96. #     pred,logit=sess.run([pred,logit], feed_dict={inputs:data, labels:true_label})
  97. #     print(pred)
  98. # ##Model has been restored. Above statement will print the saved value
  99.

2.create_model.py

  1. import tensorflow as tf
  2. import os
  3. import numpy as np
  4. import cv2
  5. import argparse
  6.  
  7.  
  8. def model(inputs):
  9.     w1=tf.Variable(tf.random_normal([3, 3, 1, 32], stddev=0.01))
  10.     w2=tf.Variable(tf.random_normal([3, 3, 32, 64], stddev=0.01))
  11.     w3=tf.Variable(tf.random_normal([3, 3, 64, 128], stddev=0.01))
  12.     w4=tf.Variable(tf.random_normal([2048, 625], stddev=0.01))
  13.     w5=tf.Variable(tf.random_normal([625, 10], stddev=0.01))
  14.  
  15.     l1_conv=tf.nn.conv2d(inputs, w1, strides=[1, 1, 1, 1], padding='SAME')
  16.     l1_relu=tf.nn.relu(l1_conv)
  17.     l1_pool=tf.nn.max_pool(l1_relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
  18.     l1_drop = tf.nn.dropout(l1_pool, 0.5)
  19.  
  20.     l2_conv=tf.nn.conv2d(l1_drop, w2, strides=[1, 1, 1, 1], padding='SAME')
  21.     l2_relu=tf.nn.relu(l2_conv)
  22.     l2_pool=tf.nn.max_pool(l2_relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
  23.     l2_drop = tf.nn.dropout(l2_pool, 0.5)
  24.  
  25.     l3_conv=tf.nn.conv2d(l2_drop, w3, strides=[1, 1, 1, 1], padding='SAME')
  26.     l3_relu=tf.nn.relu(l3_conv)
  27.     l3_pool=tf.nn.max_pool(l3_relu, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
  28.     l3_out = tf.reshape(l3_pool, [-1, 2048])
  29.     l3_drop = tf.nn.dropout(l3_out, 0.5)
  30.  
  31.     l4 = tf.nn.relu(tf.matmul(l3_drop, w4))
  32.     l4 = tf.nn.dropout(l4, 0.5)
  33.  
  34.     out = tf.matmul(l4, w5)
  35.     return out
  36.  
  37.  
  38. def model2(inputs):
  39.     w1=tf.Variable(tf.random_normal([5, 5, 1, 6], stddev=0.01))
  40.     b1 = tf.Variable(tf.truncated_normal([6]))
  41.  
  42.     w2=tf.Variable(tf.random_normal([5, 5, 6, 16], stddev=0.01))
  43.     b2 = tf.Variable(tf.truncated_normal([16]))
  44.  
  45.     w3=tf.Variable(tf.random_normal([5, 5, 16, 120], stddev=0.01))
  46.     b3 = tf.Variable(tf.truncated_normal([120]))
  47.  
  48.     w4 = tf.Variable(tf.truncated_normal([7 * 7 * 120, 80]))
  49.     b4 = tf.Variable(tf.truncated_normal([80]))
  50.  
  51.     w5 = tf.Variable(tf.truncated_normal([80, 10]))
  52.     b5 = tf.Variable(tf.truncated_normal([10]))
  53.  
  54.     l1_conv=tf.nn.conv2d(inputs, w1, strides=[1, 1, 1, 1], padding='SAME')
  55.     l1_sigmoid=tf.nn.sigmoid(l1_conv+b1)
  56.     l1_pool=tf.nn.max_pool(l1_sigmoid, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
  57.  
  58.     l2_conv=tf.nn.conv2d(l1_pool, w2, strides=[1, 1, 1, 1], padding='SAME')
  59.     l2_sigmoid=tf.nn.sigmoid(l2_conv+b2)
  60.     l2_pool=tf.nn.max_pool(l2_sigmoid, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
  61.  
  62.     l3_conv=tf.nn.conv2d(l2_pool, w3, strides=[1, 1, 1, 1], padding='SAME')
  63.     l3_sigmoid=tf.nn.sigmoid(l3_conv+b3)
  64.  
  65.     l3_out = tf.reshape(l3_sigmoid, [-1, 7*7*120])
  66.  
  67.     l4 = tf.nn.sigmoid(tf.matmul(l3_out, w4)+b4)
  68.  
  69.     out = tf.nn.softmax(tf.matmul(l4, w5) + b5)
  70.     return out
  71.  
  72.  
  73.  
  74. from datetime import datetime
  75. import time
  76. import math
  77. import collections
  78. import tensorflow as tf
  79.  
  80. slim = tf.contrib.slim
  81.  
  82. # 使用collections.namedtuple设计ResNet的Block模块
  83. # scope参数是block的名称
  84. # unit_fn是功能单元（如残差单元）
  85. # args是一个列表，如([256, 64, 1]) X 2 + [256, 64, 2]),代表两个（256, 64, 1）单元
  86. # 和一个（256, 64, 2）单元
  87. Block = collections.namedtuple("Block", ['scope', 'unit_fn', 'args'])
  88.  
  89.  
  90. # 定义下采样的方法，通过max_pool2d实现
  91. def subsample(inputs, factor, scope=None):
  92.     if factor == 1:
  93.         return inputs
  94.     else:
  95.         return slim.max_pool2d(inputs, [1, 1], stride=factor, scope=scope)
  96.  
  97.  
  98. # 定义一个创建卷积层的函数
  99. def conv2d_same(inputs, num_outputs, kernel_size, stride, scope=None):
  100.     if stride == 1:
  101.         return slim.conv2d(inputs, num_outputs, kernel_size, stride=1,
  102.                            padding='SAME', scope=scope)
  103.     else:
  104.         pad_total = kernel_size - 1
  105.         pad_beg = pad_total // 2
  106.         pad_end = pad_total - pad_beg
  107.         inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end],
  108.                                  [pad_beg, pad_end], [0, 0]])
  109.         return slim.conv2d(inputs, num_outputs, kernel_size, stride=stride,
  110.                            padding='VALID', scope=scope)
  111.  
  112.  
  113. # 定义堆叠的block函数
  114. @slim.add_arg_scope
  115. def stack_blocks_dense(net, blocks, outputs_collections=None):
  116.     for block in blocks:
  117.         with tf.variable_scope(block.scope, 'block', [net]) as sc:
  118.             for i, unit in enumerate(block.args):
  119.                 with tf.variable_scope('unit_%d' % (i + 1), values=[net]):
  120.                     unit_depth, unit_depth_bottleneck, unit_stride = unit
  121.                     net = block.unit_fn(net,
  122.                                         depth=unit_depth,
  123.                                         depth_bottleneck=unit_depth_bottleneck,
  124.                                         stride=unit_stride)
  125.             net = slim.utils.collect_named_outputs(outputs_collections, sc.name,net)
  126.  
  127.     return net
  128.  
  129.  
  130. # 用于设定默认值
  131. def resnet_arg_scope(is_training=True,
  132.                      weight_decay=0.0001,
  133.                      batch_norm_decay=0.997,
  134.                      batch_norm_epsilon=1e-5,
  135.                      batch_norm_scale=True):
  136.     batch_norm_params = {
  137.         'is_training': is_training,
  138.         'decay': batch_norm_decay,
  139.         'epsilon': batch_norm_epsilon,
  140.         'scale': batch_norm_scale,
  141.         'updates_collections': tf.GraphKeys.UPDATE_OPS,
  142.     }
  143.  
  144.     with slim.arg_scope(
  145.             [slim.conv2d],
  146.             weights_regularizer=slim.l2_regularizer(weight_decay),
  147.             weights_initializer=slim.variance_scaling_initializer(),
  148.             activation_fn=tf.nn.relu,
  149.             normalizer_fn=slim.batch_norm,
  150.             normalizer_params=batch_norm_params
  151.     ):
  152.         with slim.arg_scope([slim.batch_norm], **batch_norm_params):
  153.             with slim.arg_scope([slim.max_pool2d], padding='SAME') as arg_sc:
  154.                 return arg_sc
  155.  
  156.  
  157. # 定义残差学习单元
  158. @slim.add_arg_scope
  159. def bottleneck(inputs, depth, depth_bottleneck, stride,
  160.                outputs_collections=None, scope=None):
  161.     with tf.variable_scope(scope, 'bottleneck_v2', [inputs]) as sc:
  162.         depth_in = slim.utils.last_dimension(inputs.get_shape(), min_rank=4)
  163.         preact = slim.batch_norm(inputs, activation_fn=tf.nn.relu,
  164.                                  scope='preact')
  165.         # shortcut为直连的X
  166.         if depth == depth_in:
  167.             shortcut = subsample(inputs, stride, 'shortcut')
  168.         else:
  169.             shortcut = slim.conv2d(preact, depth, [1, 1], stride=stride,
  170.                                    normalizer_fn=None, activation_fn=None,
  171.                                    scope='shortcut')
  172.  
  173.         residual = slim.conv2d(preact, depth_bottleneck, [1, 1], stride=1,
  174.                                scope='conv1')
  175.         residual = conv2d_same(residual, depth_bottleneck, 3, stride,
  176.                                scope='conv2')
  177.         residual = slim.conv2d(residual, depth, [1, 1], stride=1,
  178.                                normalizer_fn=None, activation_fn=None,
  179.                                scope='conv3')
  180.  
  181.         # 将直连的X加到残差上，得到output
  182.         output = shortcut + residual
  183.  
  184.         return slim.utils.collect_named_outputs(outputs_collections,
  185.                                                 sc.name, output)
  186.  
  187.  
  188. # 定义ResNet的主函数
  189. def resnet_v2(inputs,
  190.               blocks,
  191.               num_classes=None,
  192.               global_pool=True,
  193.               include_root_block=True,
  194.               reuse=None,
  195.               scope=None):
  196.     with tf.variable_scope(scope, 'resnet_v2', [inputs], reuse=reuse) as sc:
  197.         end_points_collection = sc.original_name_scope + '_end_points'
  198.         
  199.         with slim.arg_scope([slim.conv2d, bottleneck,stack_blocks_dense],outputs_collections=end_points_collection):
  200.             net = inputs
  201.             
  202.             if include_root_block:
  203.                 with slim.arg_scope([slim.conv2d], activation_fn=None,normalizer_fn=None):
  204.                     net = conv2d_same(net, 64, 7, stride=2, scope='conv1')
  205.                 net = slim.max_pool2d(net, [3, 3], stride=2, scope='pool1')
  206.  
  207.             net = stack_blocks_dense(net, blocks)
  208.             net = slim.batch_norm(net, activation_fn=tf.nn.relu, scope='postnorm')
  209.             
  210.             if global_pool:
  211.                 net = tf.reduce_mean(net, [1, 2], name='pool5', keep_dims=True)
  212.  
  213.             if num_classes is not None:
  214.                 net = slim.conv2d(net, num_classes, [1, 1], activation_fn=None,normalizer_fn=None, scope='logits')
  215.  
  216.             end_points = slim.utils.convert_collection_to_dict(end_points_collection)
  217.  
  218.             if num_classes is not None:
  219.                 end_points['predictions'] = slim.softmax(net, scope='predictions')
  220.             return net, end_points
  221.  
  222.  
  223. # 定义50层的ResNet
  224. def resnet_v2_50(inputs,
  225.                   num_classes=None,
  226.                   global_pool=True,
  227.                   reuse=None,
  228.                   scope='resnet_v2_50'):
  229.     blocks = [
  230.         Block('block1', bottleneck, [(256, 64, 1)] * 2 + [(256, 64, 2)]),
  231.         Block('block2', bottleneck, [(512, 128, 1)] * 3 + [(512, 128, 2)]),
  232.         Block('block3', bottleneck, [(1024, 256, 1)] * 5 + [(1024, 256, 2)]),
  233.         Block('block4', bottleneck, [(2048, 512, 1)] * 3)]
  234.     return resnet_v2(inputs, blocks, num_classes, global_pool,include_root_block=True, reuse=reuse, scope=scope)

3.data_loader.py

  1. import cv2
  2. import os
  3. import numpy as np
  4. import random
  5.  
  6. # f1=open('dataset/train_data.txt','w+')
  7. # path='dataset/'
  8. # for file in os.listdir(path):
  9. #     if file.endswith('png'):
  10. #         line=path+file+'  1'+'\n'
  11. #         f1.write(line)
  12. #         print(line)
  13.  
  14.  
  15.  
  16. def one_hot(data, num_classes):
  17.   return np.squeeze(np.eye(num_classes)[data.reshape(-1)])
  18.  
  19.  
  20. def get_data_list(path):
  21.     f1=open(path,'r')
  22.     lines=f1.readlines()
  23.  
  24.     img_list=[]
  25.     label_list=[]
  26.     for line in lines:
  27.         label=int(line.strip().split("  ")[1])
  28.         label=one_hot(np.array(label),10)
  29.         label_list.append(label)
  30.  
  31.         file_name=line.strip().split("  ")[0]
  32.         img=cv2.imread(file_name, 0)
  33.         img=np.reshape(img,[28,28,1])
  34.         # print(img.shape)
  35.         img_list.append(img)
  36.         
  37.     return img_list, label_list
  38.  
  39.  
  40. def get_data(img_list, label_list, batch_size):
  41.     lens=len(label_list)
  42.     random_nums=random.sample(range(lens),lens)
  43.     nums=random_nums[0:batch_size]
  44.     # print(nums)
  45.  
  46.     data=[]
  47.     label=[]
  48.     for index in nums:
  49.         data.append(img_list[index])
  50.         label.append(label_list[index])
  51.         
  52.     return np.array(data), np.array(label)
  53.  
  54.  
  55.  
  56. # batch_size=1
  57. # path="dataset/mnist/train/train_data.txt"
  58.  
  59. # img_list, label_list=get_data_list(path)
  60. # print(len(img_list), len(label_list))
  61.  
  62. # data, label=get_data(img_list, label_list, batch_size)
  63. # print(type(data[0]),label)

4.查看ckpt网络.py

  1. from tensorflow.python import pywrap_tensorflow
  2. import os
  3. import tensorflow as tf
  4. from tensorflow.python.platform import gfile
  5.  
  6. # checkpoint_path = os.path.join('model/model.ckpt')
  7. # reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
  8. # var_to_shape_map = reader.get_variable_to_shape_map()
  9. # for key in var_to_shape_map:
  10. #     print('tensor_name: ', key)
  11.  
  12. ckpt_path = os.path.join('model/model')
  13. saver = tf.train.import_meta_graph(ckpt_path+'.meta',clear_devices=True)
  14. graph = tf.get_default_graph()
  15. with tf.Session( graph=graph) as sess:
  16.     sess.run(tf.global_variables_initializer()) 
  17.     saver.restore(sess,ckpt_path)
  18.  
  19.     # tensorboard --logdir=d:/log  --host=127.0.0.1
  20.     init = tf.initialize_all_variables()
  21.     merged = tf.summary.merge_all()
  22.     writer = tf.summary.FileWriter("d:/log/",sess.graph) #目录结构尽量简单，复杂了容易出现找不到文件，原因不清楚
  23.     sess.run(init)

5.ckpt2pb.py

  1. import tensorflow as tf
  2. from tensorflow.python.framework import graph_util
  3. from tensorflow.python import pywrap_tensorflow
  4.  
  5.  
  6. def freeze_graph(cpkt_path, pb_path):
  7.     checkpoint = tf.train.get_checkpoint_state("model") #检查目录下ckpt文件状态是否可用
  8.     cpkt_path2 = checkpoint.model_checkpoint_path #得ckpt文件路径
  9.     print("gg:",checkpoint,cpkt_path2)
  10.  
  11.     # 指定输出的节点名称,该节点名称必须是原模型中存在的节点
  12.     output_node_names = "resnet_v2_50/logits/BiasAdd"
  13.     saver = tf.train.import_meta_graph(cpkt_path + '.meta', clear_devices=True)
  14.     graph = tf.get_default_graph()
  15.     input_graph_def = graph.as_graph_def()
  16.  
  17.     # feature_data_list = input_graph_def.get_operation_by_name('resnet_v2_50/conv1').outputs[0]
  18.     # input_image=tf.placeholder(None,28,28,1)
  19.  
  20.     with tf.Session() as sess:
  21.         saver.restore(sess, cpkt_path)  # 恢复图并得到数据
  22.  
  23.         pb_path_def = graph_util.convert_variables_to_constants(  # 模型持久化，将变量值固定
  24.             sess=sess,
  25.             input_graph_def=input_graph_def,  # 等于:sess.graph_def
  26.             output_node_names=output_node_names.split(","))  # 如果有多个输出节点，以逗号隔开
  27.         # print(pb_path_def)
  28.  
  29.  
  30.         with tf.gfile.GFile(pb_path, 'wb') as fgraph:
  31.             fgraph.write(pb_path_def.SerializeToString())
  32.         # with tf.io.gfile.GFile(pb_path, "wb") as f:  # 保存模型
  33.         #     f.write(pb_path_def.SerializeToString())  # 序列化输出
  34.         print("%d ops in the final graph." % len(pb_path_def.node))  # 得到当前图有几个操作节点
  35.  
  36. if __name__ == '__main__':
  37.     # 输入路径(cpkt)
  38.     cpkt_path = 'model/model'
  39.     # 输出路径(pb模型)pb_path_def
  40.     pb_path = "model/test.pb"
  41.     # 模型转换
  42.     freeze_graph(cpkt_path, pb_path)
  43.  
  44.     # # 查看节点名称：
  45.     # reader = pywrap_tensorflow.NewCheckpointReader(cpkt_path)
  46.     # var_to_shape_map = reader.get_variable_to_shape_map()
  47.     # for key in var_to_shape_map:
  48.     #     print("tensor_name: ", key)
  49.  
  50.     # # 查看某个指定节点的权重
  51.     # reader = pywrap_tensorflow.NewCheckpointReader(cpkt_path)
  52.     # var_to_shape_map = reader.get_variable_to_shape_map()
  53.     # w0 = reader.get_tensor("finetune/dense_1/bias")
  54.     # print(w0.shape, type(w0))
  55.     # print(w0[0])
  56.     
  57.  
  58.     # with tf.Session() as sess:
  59.     #     # 加载模型定义的graph
  60.     #     saver = tf.train.import_meta_graph('model/model.meta')
  61.     #     # 方式一:加载指定文件夹下最近保存的一个模型的数据
  62.     #     saver.restore(sess, tf.train.latest_checkpoint('model/'))
  63.     #     # 方式二:指定具体某个数据，需要注意的是，指定的文件不要包含后缀
  64.     #     # saver.restore(sess, os.path.join(path, 'model.ckpt-1000'))
  65.  
  66.     #     # 查看模型中的trainable variables
  67.     #     tvs = [v for v in tf.trainable_variables()]
  68.     #     for v in tvs:
  69.     #         print(v.name)
  70.     #         # print(sess.run(v))
  71.  
  72.     #     # # 查看模型中的所有tensor或者operations
  73.     #     # gv = [v for v in tf.global_variables()]
  74.     #     # for v in gv:
  75.     #     #     print(v.name)
  76.  
  77.     #     # # 获得几乎所有的operations相关的tensor
  78.     #     # ops = [o for o in sess.graph.get_operations()]
  79.     #     # for o in ops:
  80.     #     #     print(o.name)

6.pb2pbtxt.py   (pb和pbtxt互转，修改input_shape)

修改pbtxt文件，把动态shape改成静态的，再转回pb模型

  1. import tensorflow as tf
  2. from tensorflow.python.platform import gfile
  3. from google.protobuf import text_format
  4.  
  5.  
  6. def convert_pb_to_pbtxt(root_path, pb_path, pbtxt_path):
  7.     with gfile.FastGFile(root_path+pb_path, 'rb') as f:
  8.         graph_def = tf.GraphDef()
  9.  
  10.         graph_def.ParseFromString(f.read())
  11.  
  12.         tf.import_graph_def(graph_def, name='')
  13.  
  14.         tf.train.write_graph(graph_def, root_path, pbtxt_path, as_text=True)
  15.     return
  16.  
  17.  
  18. def convert_pbtxt_to_pb(root_path, pb_path, pbtxt_path):
  19.     with tf.gfile.FastGFile(root_path+pbtxt_path, 'r') as f:
  20.         graph_def = tf.GraphDef()
  21.         file_content = f.read()
  22.  
  23.         # Merges the human-readable string in `file_content` into `graph_def`.
  24.         text_format.Merge(file_content, graph_def)
  25.         tf.train.write_graph(graph_def, root_path, pb_path, as_text=False)
  26.     return
  27.  
  28. if __name__ == '__main__':
  29.     # 模型路径
  30.     root_path = "model/"
  31.     pb_path = "test.pb"
  32.     pbtxt_path = "test.pbtxt"
  33.  
  34.     # 模型转换
  35.     convert_pb_to_pbtxt(root_path, pb_path, pbtxt_path)
  36.     # convert_pbtxt_to_pb(root_path, pb_path, pbtxt_path)

7.test_pb.py

  1. import tensorflow as tf
  2. from tensorflow.python.framework import graph_util
  3. from tensorflow.python import pywrap_tensorflow
  4. import numpy as np
  5.  
  6. def recognize(jpg_path, pb_file_path):
  7.     with tf.Graph().as_default():
  8.         output_graph_def = tf.GraphDef()
  9.  
  10.         with open(pb_file_path, "rb") as f:#主要步骤即为以下标出的几步,1、2步即为读取图
  11.             output_graph_def.ParseFromString(f.read())# 1.将模型文件解析为二进制放进graph_def对象
  12.             _ = tf.import_graph_def(output_graph_def, name="")# 2.import到当前图
  13.  
  14.         with tf.Session() as sess:
  15.             init = tf.global_variables_initializer()
  16.             sess.run(init)
  17.  
  18.             graph = tf.get_default_graph()# 3.获得当前图
  19.             
  20.             # # 4.get_tensor_by_name获取需要的节点
  21.             # x = graph.get_tensor_by_name("IteratorGetNext_1:0")
  22.             # y_out = graph.get_tensor_by_name("resnet_v1_50_1/predictions/Softmax:0")
  23.  
  24.             x = graph.get_tensor_by_name("input_images:0")
  25.             y_out = graph.get_tensor_by_name("resnet_v2_50/logits/BiasAdd:0")
  26.             
  27.             img=np.random.normal(size=(1, 28, 28, 1))
  28.             # img=cv2.imread(jpg_path)
  29.             # img=cv2.resize(img, (224, 224))
  30.             # img=np.reshape(img,(1,224,224,3))
  31.             # print(img.shape)
  32.             
  33.             #执行
  34.             output = sess.run(y_out, feed_dict={x:img})
  35.             pred=np.argmax(output[0][0], axis=1)
  36.             print("预测结果：", output.shape, output, "预测label:", pred)
  37.  
  38.             # prediction_labels = np.argmax(test_y_out, axis=2)
  39.             # print(prediction_labels.shape, prediction_labels)
  40. recognize("dataset/00000.PNG", "model/test.pb")

声明：本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：【wpsshop博客】
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号