项目背景

EfficientNet系列博客与ResNet迁移学习的任务一样，均是进行二分类的迁移学习，唯一差别在于网络结构的不同。其它相关的内容，参考下面链接：ResNet迁移学习（一）—— 主程序结构，部分内容如下图所示：

代码框架

下面我们主要介绍train.py的程序框架和功能模块接口。程序框架是构建工程的首要步骤，而功能模块是完善框架的内容。这样在构建工程的时候，思路会比较清晰，不至于无从下手。

1. 数据接口

 tf.set_random_seed(-1)

# ****************************************************************** #
#                   1. Python多线程数据读取与数据预处理                 #
# ****************************************************************** # 
train_queue = train_set_queue() 
valid_queue = valid_set_queue()

2. 网络结构接口

    # ****************************************************************** #
    #                   2. 构建静态图                                     #
    # ****************************************************************** #
    model = Model()
    print('variables: ', model.first_stage_trainable_var_list)

3. 开启会话，执行训练

第一步：如果要加载预训练模型，则需要首先恢复其相应的权重

# ****************************************************************** #
#                   3. 恢复权重, 开启会话, 执行静态图                   #
# ****************************************************************** #
    init = tf.global_variables_initializer()
    config_proto = tf.ConfigProto()
    config_proto.gpu_options.allow_growth = True
    with tf.Session(config=config_proto) as sess:
        sess.run(init)
        try:
            print('=> Restoring weights from: %s ... ' % model.pre_trained_weight)
            model.loader.restore(sess, model.pre_trained_weight)
            print('成功加载模型')
        except:
            print('Load Pretrained Weight Failed !!!\n')
            print('=> %s does not exist !!!' % model.pre_trained_weight)
            print('=> Now it starts to train from scratch ...')
            cfg.Train.First_Stage_Epochs = 0  # 重头开始训练网络, 不需要预训练权重
        else:
            print('\nLoad Model Success !!!\n')

第二步：设置训练过程中的保存文件目录，将要收集的变量保存文件，可以用Tensorboard查看

        start_epoch = 0

        # 训练日志保存目录
        save_path = cfg.train.log
        if os.path.exists(save_path):
            print('目录非空, 清空目录')
            shutil.rmtree(save_path)
        train_writer = tf.summary.FileWriter(save_path + 'train', sess.graph)
        valid_writer = tf.summary.FileWriter(save_path + 'valid')

第三步：根据设置的Epoch，循环训练网络

        for epoch in range(start_epoch + 1, 1 + model.first_stage_epoch + model.second_stage_epoch):
            if epoch < model.first_stage_epoch:
                train_op = model.train_op_with_frozen_variables
            else:
                train_op = model.train_op_with_all_variables

            # Train Process
            for step in range(0, cfg.train.train_num // model.batch_size):
                # 每次从队列中取出一个batch的数据
                _, _, image, label = train_queue.get()

                # 统计每个batch中标签的二分类分布情况
                # result = Counter(label)
                train_params = [train_op,
                                model.loss,
                                model.accuracy,
                                model.merged,
                                model.learn_rate,
                                model.global_step]

                train_dict = {
     model.inputs: image,
                              model.label_c: label,
                              model.trainable: True}

                _, train_loss, train_accuracy, merge_train, lr_, global_step_ = sess.run(train_params,
                                                                                         feed_dict=train_dict)

                train_writer.add_summary(merge_train, global_step_)
                print('iter:%2d/%d || train loss:%.4f || train accuracy:%.4f || lr:%g ' % (step,
                                                                                           epoch,
                                                                                           train_loss,
                                                                                           train_accuracy,
                                                                                           lr_))
                # print('阴性样本：%d || 阳性样本：%d ' % (result[0], result[1]))

            # Valid Process
            valid = []
            label_ = []
            predict_ = []
            for valid_step in range(cfg.train.valid_num // model.batch_size):
                name, _, valid_image, valid_label = valid_queue.get()

                valid_params = [model.accuracy,
                                model.loss,
                                model.merged,
                                tf.argmax(model.logits, 1),
                                model.label_c]

                valid_dict = {
     model.inputs: valid_image,
                              model.label_c: valid_label,
                              model.trainable: False}

                valid_accuracy, valid_loss, merge_valid, predict, label_c = sess.run(valid_params, feed_dict=valid_dict)

                valid_writer.add_summary(merge_valid, global_step_)
                print('==================================================================================')
                print('step=%2d/%d || valid accuracy=%g || loss=%.4f ' % (valid_step,
                                                                          epoch,
                                                                          valid_accuracy,
                                                                          valid_loss))
                print('==================================================================================')
                print('标签值：', label_c)
                print('预测值：', predict)

                label_.append(label_c)
                predict_.append(predict)
                valid.append(valid_accuracy)

            # 计算所有验证集的混淆矩阵
            print('mean valid accuracy: ', np.mean(valid))
            confusion_matrix = tf.confusion_matrix(np.hstack(label_), np.hstack(predict_), num_classes=2)
            confusion_matrix_ = sess.run(confusion_matrix)

            TN = confusion_matrix_[0][0]
            FP = confusion_matrix_[0][1]
            FN = confusion_matrix_[1][0]
            TP = confusion_matrix_[1][1]

            acc = (TP + TN) / (TP + TN + FP + FN)
            sensitive = TP / (TP + FN)
            specify = TN / (TN + FP)

            print('混淆矩阵\n', confusion_matrix_)
            print('准确度, 灵敏度, 特异度: ', acc, sensitive, specify)

            # 保存每个epoch的模型
            if epoch >= 79:
                exit()
            model.saver.save(sess, cfg.train.save_model, global_step=epoch)

完整代码

from Model import *
from DataLoader import *
from config import cfg
from collections import Counter
import shutil
os.environ['CUDA_VISIBLE_DEVICES'] = '1'


def train():
    tf.set_random_seed(-1)

    # ****************************************************************** #
    #                   1. Python多线程数据读取与数据预处理                 #
    # ****************************************************************** #
    train_queue = train_set_queue()
    valid_queue = valid_set_queue()
    # for i in range(1000000):
    #     name1, _, image, label = train_queue.get()
    #     name2, _, valid_image, valid_label = valid_queue.get()
    #     print(name1)
    #     print(name2)
    #     print(len(name1))
    #     print(len(name2))
    #     print('==============================================')

    # ****************************************************************** #
    #                   2. 构建静态图                                     #
    # ****************************************************************** #
    model = Model()
    print('variables: ', model.first_stage_trainable_var_list)

    # ****************************************************************** #
    #                   3. 恢复权重, 开启会话, 执行静态图                   #
    # ****************************************************************** #
    init = tf.global_variables_initializer()
    config_proto = tf.ConfigProto()
    config_proto.gpu_options.allow_growth = True
    with tf.Session(config=config_proto) as sess:
        sess.run(init)
        try:
            print('=> Restoring weights from: %s ... ' % model.pre_trained_weight)
            model.loader.restore(sess, model.pre_trained_weight)
            print('成功加载模型')
        except:
            print('Load Pretrained Weight Failed !!!\n')
            print('=> %s does not exist !!!' % model.pre_trained_weight)
            print('=> Now it starts to train from scratch ...')
            cfg.Train.First_Stage_Epochs = 0  # 重头开始训练网络, 不需要预训练权重
        else:
            print('\nLoad Model Success !!!\n')

        start_epoch = 0

        # 训练日志保存目录
        save_path = cfg.train.log
        if os.path.exists(save_path):
            print('目录非空, 清空目录')
            shutil.rmtree(save_path)
        train_writer = tf.summary.FileWriter(save_path + 'train', sess.graph)
        valid_writer = tf.summary.FileWriter(save_path + 'valid')

        for epoch in range(start_epoch + 1, 1 + model.first_stage_epoch + model.second_stage_epoch):
            if epoch < model.first_stage_epoch:
                train_op = model.train_op_with_frozen_variables
            else:
                train_op = model.train_op_with_all_variables

            # Train Process
            for step in range(0, cfg.train.train_num // model.batch_size):
                # 每次从队列中取出一个batch的数据
                _, _, image, label = train_queue.get()

                # 统计每个batch中标签的二分类分布情况
                # result = Counter(label)
                train_params = [train_op,
                                model.loss,
                                model.accuracy,
                                model.merged,
                                model.learn_rate,
                                model.global_step]

                train_dict = {
    model.inputs: image,
                              model.label_c: label,
                              model.trainable: True}

                _, train_loss, train_accuracy, merge_train, lr_, global_step_ = sess.run(train_params,
                                                                                         feed_dict=train_dict)

                train_writer.add_summary(merge_train, global_step_)
                print('iter:%2d/%d || train loss:%.4f || train accuracy:%.4f || lr:%g ' % (step,
                                                                                           epoch,
                                                                                           train_loss,
                                                                                           train_accuracy,
                                                                                           lr_))
                # print('阴性样本：%d || 阳性样本：%d ' % (result[0], result[1]))

            # Valid Process
            valid = []
            label_ = []
            predict_ = []
            for valid_step in range(cfg.train.valid_num // model.batch_size):
                name, _, valid_image, valid_label = valid_queue.get()

                valid_params = [model.accuracy,
                                model.loss,
                                model.merged,
                                tf.argmax(model.logits, 1),
                                model.label_c]

                valid_dict = {
    model.inputs: valid_image,
                              model.label_c: valid_label,
                              model.trainable: False}

                valid_accuracy, valid_loss, merge_valid, predict, label_c = sess.run(valid_params, feed_dict=valid_dict)

                valid_writer.add_summary(merge_valid, global_step_)
                print('==================================================================================')
                print('step=%2d/%d || valid accuracy=%g || loss=%.4f ' % (valid_step,
                                                                          epoch,
                                                                          valid_accuracy,
                                                                          valid_loss))
                print('==================================================================================')
                print('标签值：', label_c)
                print('预测值：', predict)

                label_.append(label_c)
                predict_.append(predict)
                valid.append(valid_accuracy)

            # 计算所有验证集的混淆矩阵
            print('mean valid accuracy: ', np.mean(valid))
            confusion_matrix = tf.confusion_matrix(np.hstack(label_), np.hstack(predict_), num_classes=2)
            confusion_matrix_ = sess.run(confusion_matrix)

            TN = confusion_matrix_[0][0]
            FP = confusion_matrix_[0][1]
            FN = confusion_matrix_[1][0]
            TP = confusion_matrix_[1][1]

            acc = (TP + TN) / (TP + TN + FP + FN)
            sensitive = TP / (TP + FN)
            specify = TN / (TN + FP)

            print('混淆矩阵\n', confusion_matrix_)
            print('准确度, 灵敏度, 特异度: ', acc, sensitive, specify)

            # 保存每个epoch的模型
            if epoch >= 79:
                exit()
            model.saver.save(sess, cfg.train.save_model, global_step=epoch)


if __name__ == '__main__':
    train()