【Tensorflow】Inception-v3 迁移学习

数据集下载：经典的cat数据
inception-v3下载：classify_image_graph_def.pb

一、注意事项

• 创建sess时启动inception-v3： tf.Session(graph=graph)；
• 最后一层全连接层的输入使用tf.placeholder_with_default(）占位，而不是tf.placeholder（）；
• 训练时记得保存labels.txt，用于后续使用remodel.pb；（主要是label值和顺序）
• 在使用remodel.pb时，注意导入tensor的名称书写正确性；（可通过查看tensorboard中graph进行判断）

二、使用Inception-v3进行retrain

1、指定路径和参数
路径：

• 输入：数据集（注意文件夹格式）、inception-v3存放路径、
• 输出：inception-v3模型输出值、tensorboard训练日志、重训练后的新模型、数据集labels保存

参数：

• 验证集\测试集占比、迭代次数、batch大小、inception-v3模型output的节点数、学习率

import numpy as np
import tensorflow as tf
import os
import glob
import random
from tensorflow.python.framework import graph_util
import shutil

path = r'E:\cat' #主文件夹
datapath = os.path.join(path, 'dataset') #图像数据路径
bottleneckpath = os.path.join(path, 'bottleneck') #inception-v3输出结果缓存路径
premodel_path = os.path.join(path, 'inception-2015-12-05', 'classify_image_graph_def.pb') #inception-v3模型路径
log_path = os.path.join(path, 'log') #tensorboard训练日志保存路径
remodel_path = os.path.join(path, 'remodel', 'remodel.pb') #重训练后的模型保存路径
label_path = os.path.join(path, 'labels.txt') #数据集labels保存路径（用于retrain.pb测试）

valpct = 10 #验证集占比
testpct = 10 #测试集占比

max_steps = 6000 #迭代次数
batch = 100 #batch_size
bottleneck_size = 2048 #inception-v3模型output的节点数（固定）
learn_rate = 0.01 #学习率

tf.reset_default_graph()

2、导入inception-v3，构建最后一层全连接层，对新数据训练验证测试

#读数据：获取图片路径
def get_img_path():
    sub_dirs = [i[0] for i in os.walk(datapath)]
    is_root_file = True
    results = {}
    for sub_dir in sub_dirs:
        if is_root_file:
            is_root_file = False
            continue
        label = os.path.basename(sub_dir).lower()

        imgs = []
        imgtypes = ['jpg', 'jpeg', 'JPG', 'JPEG']
        for imgtype in imgtypes:
           imgs.extend( glob.glob(os.path.join(sub_dir, '*.'+imgtype))) 

        train = []
        val = []
        test =[]
        for img in imgs:
            channce = np.random.randint(100)
            if channce < valpct:
                test.append(img)
            elif channce < testpct + valpct:
                val.append(img)
            else:
                train.append(img)
        results[label] = {
                'train': train,
                'test': test,
                'val': val}
    return results


# 跑inception-v3的output数据
def run_bottleneck_tensor(sess, bottleneck_tensor, jpeg_data_tensor, image_data):
    bottleneck_value = sess.run(bottleneck_tensor, {jpeg_data_tensor: image_data})
    bottleneck_value = np.squeeze(bottleneck_value)
    return bottleneck_value


# 获取inception模型数据
def get_or_create_bottleneck_values(sess, bottleneck_tensor, jpeg_data_tensor, label_name, img_path):
    bottleneck_label_path = os.path.join(bottleneckpath, label_name)
    if not os.path.exists(bottleneck_label_path):
        os.makedirs(bottleneck_label_path)
    img_name = os.path.basename(img_path)
    bottleneck_img_path = os.path.join(bottleneck_label_path, img_name+'.txt')

    if not os.path.exists(bottleneck_img_path):
        img_data = tf.gfile.FastGFile(img_path, 'rb').read()
        bottleneck_value = run_bottleneck_tensor(sess, bottleneck_tensor, jpeg_data_tensor, img_data)
        with open(bottleneck_img_path, 'w') as f:
            f.write(','.join(str(i) for i in bottleneck_value))
    else:
        with open(bottleneck_img_path, 'r') as f:
            bottleneck_value = [ float(i) for i in f.read().split(',')]
    return bottleneck_value


# 随机获取train/valiation的batch数据
def get_batch(sess, batch, img_lists, n_classes, bottleneck_tensor, jpeg_data_tensor, ops):
    bottleneck_values = [] #inception-v3的output值
    y = [] #图片数据的标签值（onehot格式）
    for i in range(batch):
        label_index = random.randrange(n_classes)
        label_name = list(img_lists.keys())[label_index]
        label = np.zeros(n_classes)
        label[label_index] = 1.0
        y.append(label)

        samples = len(img_lists[label_name][ops])
        img_index = random.randrange(samples)
        img_path =  img_lists[label_name][ops][img_index]
        bottleneck_value = get_or_create_bottleneck_values(sess, bottleneck_tensor, jpeg_data_tensor, label_name, img_path)
        bottleneck_values.append(bottleneck_value)

    return bottleneck_values, y


# 获取所有test数据   
def get_test(sess, img_lists, n_classes, bottleneck_tensor, jpeg_data_tensor):
    bottleneck_values = []
    y = []
    for label_index in range(n_classes):
        label_name = list(img_lists.keys())[label_index]
        label = np.zeros(n_classes)
        label[label_index] = 1.0
        n_img = len(img_lists[label_name]['test'])
        for img_index in range(n_img):
            img_path =  img_lists[label_name]['test'][img_index]
            bottleneck_value = get_or_create_bottleneck_values(sess, bottleneck_tensor, jpeg_data_tensor, label_name, img_path)
            bottleneck_values.append(bottleneck_value)
            y.append(label)
    return bottleneck_values, y


def main(_):  

    img_lists = get_img_path()
    n_classes = len(img_lists.keys())

    # 保存labels数据：主要是label值和label顺序，用于后续对模型的使用
    with open(label_path, 'w') as f:
        f.write('\n'.join(img_lists.keys())) 

    # 清空log文件夹
    if not os.path.exists(log_path):
        os.makedirs(log_path)
    else:
        shutil.rmtree(log_path)
        os.makedirs(log_path)  


    def variable_summary(var):
        with tf.name_scope('summaries'):
            mean = tf.reduce_mean(var)
            tf.summary.scalar('mean', mean)
            with tf.name_scope('stddev'):
                std = tf.sqrt(tf.reduce_mean(tf.square( var - mean )))
            tf.summary.scalar('std', std)
            tf.summary.scalar('max', tf.reduce_max(var))
            tf.summary.scalar('min', tf.reduce_min(var))
            tf.summary.histogram('histogram', var)


    # 导入inception-v3模型
    with tf.Graph().as_default() as graph:
        with tf.gfile.FastGFile(premodel_path, 'rb') as f:
            inception = tf.GraphDef()
            inception.ParseFromString(f.read())
            # import两个tensor:'pool_3/_reshape:0'(output), 'DecodeJpeg/contents:0'(input)
            bottleneck_tensor, jpeg_data_tensor = tf.import_graph_def(inception, return_elements=['pool_3/_reshape:0', 'DecodeJpeg/contents:0'])   

    with tf.Session(graph=graph) as sess: 
        # 创建最后一层 
        with tf.name_scope('final_layer'):
            with tf.name_scope('input'):
                bottleneck_input = tf.placeholder_with_default(
                        bottleneck_tensor, 
                        [None, bottleneck_size], 
                        name='bottleneck_input')
                y_input = tf.placeholder(tf.float32, [None, n_classes], name='y_input')
            with tf.name_scope('weights'):
                initial_weights = tf.truncated_normal([bottleneck_size, n_classes], stddev=0.001)
                layer_weights = tf.Variable(initial_weights, name='final_layer_weights')
                variable_summary(layer_weights)
            with tf.name_scope('biases'):
                initial_biases = tf.zeros([n_classes])
                layer_biases = tf.Variable(initial_biases, name='layer_biases')
                variable_summary(layer_biases)
            with tf.name_scope('wx_plus_b'):
                logits = tf.matmul(bottleneck_input, layer_weights) + layer_biases
                tf.summary.histogram('pre_activations', logits)
            final_tensor = tf.nn.softmax(logits, name='activations')
            tf.summary.histogram('activations', final_tensor)

        # 定义损失函数
        with tf.name_scope('cross_entropy'):
            cross_entory = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=final_tensor, labels=y_input))
        tf.summary.scalar('cross_entory', cross_entory)

        with tf.name_scope('train'):
            train_step = tf.train.GradientDescentOptimizer(learn_rate).minimize(cross_entory)

        # 计算准确率
        with tf.name_scope('accuracy'):
            predict = tf.argmax(final_tensor, 1, name='predict')
            accuracy = tf.reduce_mean(tf.cast(tf.equal(predict, tf.argmax(y_input, 1)), tf.float32))
        tf.summary.scalar('accuracy', accuracy)

        megered = tf.summary.merge_all()
        train_log = tf.summary.FileWriter(log_path+'/train', sess.graph)
        val_log = tf.summary.FileWriter(log_path+'/val', sess.graph)
        test_log = tf.summary.FileWriter(log_path+'/test', sess.graph)

        sess.run(tf.global_variables_initializer())
        bottlenecks, y = get_batch(sess, batch, img_lists, n_classes, bottleneck_tensor, jpeg_data_tensor, 'train')

        # 训练
        for i in range(max_steps):
            bottlenecks, y = get_batch(sess, batch, img_lists, n_classes, bottleneck_tensor, jpeg_data_tensor, 'train')
            summary_train, _ = sess.run([megered, train_step], feed_dict={bottleneck_input: bottlenecks, y_input: y})
            train_log.add_summary(summary_train, i)
            if i % 100 == 0 or i + 1 == max_steps:
                bottlenecks, y = get_batch(sess, batch, img_lists, n_classes, bottleneck_tensor, jpeg_data_tensor, 'val')
                summary_val, acc = sess.run([megered, accuracy], feed_dict={bottleneck_input: bottlenecks, y_input: y}) 
                val_log.add_summary(summary_val, i)
                print('step %d:  Validation accuracy on random sampled 100 examples = %.1f%%' % (i, acc*100))

        # 测试
        bottlenecks, y = get_test(sess, img_lists, n_classes, bottleneck_tensor, jpeg_data_tensor)
        summary_test, acc = sess.run([megered, accuracy], feed_dict={bottleneck_input: bottlenecks, y_input: y}) 
        test_log.add_summary(summary_test)
        print('Final test accuracy =  %.1f%%' % (acc*100))

        # 保存训练好的模型
        out_graph_def = graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=['accuracy/predict'])
        with tf.gfile.FastGFile(remodel_path, mode='wb') as f:
            f.write(out_graph_def.SerializeToString())

if __name__ == '__main__':
    tf.app.run()

三、使用remodel.pb进行分类

import tensorflow as tf
import os


# 各种路径
base_path = r'E:\cat'
model_path = os.path.join(base_path, r'remodel\remodel.pb')
img_path =  os.path.join(base_path, 'Testdata')
label_path = os.path.join(base_path, 'labels.txt')


# 导入模型
model_file = tf.gfile.FastGFile(model_path, 'rb').read()
retrain = tf.GraphDef()
retrain.ParseFromString(model_file)
tf.import_graph_def(retrain, name='')


# 创建对话
sess = tf.Session()
input_tensor = sess.graph.get_tensor_by_name('import/DecodeJpeg/contents:0')
softmax_tensor = sess.graph.get_tensor_by_name('final_layer/activations:0')

# 遍历测试数据
lable_data = [row.rstrip() for row in tf.gfile.GFile(label_path)]
img_lists = os.listdir(img_path)
for img in img_lists:
    print('Picture: %s' % img)
    img_data = tf.gfile.FastGFile(os.path.join(img_path, img), 'rb').read()
    softmax_values = sess.run(softmax_tensor, feed_dict={input_tensor: img_data})
    top_k = softmax_values[0].argsort()[-len(softmax_values[0]):][::-1]
    for k in top_k:
        score = softmax_values[0][k]
        label = lable_data[k]
        print('\t%s (score = %.5f)' % (label, score))