迁移学习是在已经训练好的模型上，通过微调最后几层神经网络，从而达到快速训练自定义模型的一种技术
本例程展示如何基于tensorflow通过迁移学习快速训练出自定义的分类模型

数据准备

在主机目录/home/pi/Lepi_Data/ros/transfer_learning下新建文件夹”分类测试”，然后在”分类测试”目录下新建不同的分类目录，比如 “分类1”，”分类2”，向每个分类添加不少于20张对应分类的图片，然后可以运行下面这段示例程序

  1. #!coding:utf-8
  2. import tensorflow as tf
  3. from tensorflow import keras
  4. import os
  5. from tensorflow.keras.models import load_model
  6. import tensorflow.python.keras.backend as K
  7. from tensorflow.keras.preprocessing.image import img_to_array
  8. from sklearn.model_selection import train_test_split
  9. from tensorflow.keras.models import Model
  10. from tensorflow.keras.layers import GlobalAveragePooling2D, Dense
  11. from tensorflow.keras.optimizers import SGD
  12. import cv2
  13. import numpy as np
  16. def get_labels(data_dir):
  17.     """
  18.     加载标签
  19.     """
  20.     labels = []
  21.     files = os.listdir(data_dir)
  22.     files.sort()
  23.     for file in files:
  24.         if not os.path.isdir(os.path.join(data_dir, file)):
  25.             continue
  26.         catname_full, _ = os.path.splitext(file)
  27.         # catname = catname_full.split('_')[-1]
  28.         # if len(os.listdir(os.path.join(data_dir,file)))>=1:
  29.         labels.append(catname_full)
  30.     print(labels)
  31.     labels.sort()
  32.     return labels
  35. def load_img_from_dir(data_dir, target_size=(112, 112), max_num=100):  # From Directory
  36.     """
  37.     从目录加载图片并缩放至指定尺寸
  38.     """
  39.     x_load = []
  40.     y_load = []
  41.     labels = get_labels(data_dir)
  42.     dirs = labels
  43.     for cat in dirs:  # load directory
  44.         files_dir = os.path.join(data_dir, cat)
  45.         files = os.listdir(files_dir)
  46.         for file in files[:max_num]:
  47.             file_path = os.path.join(files_dir, file)
  48.             try:
  49.                 cv_img = cv2.imread(file_path)
  50.                 # cv_img=cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1)
  51.                 x = img_to_array(cv2.resize(cv_img, target_size))
  52.                 # x = K.expand_dims(x, axis=0)
  53.                 # x = preprocess_input(x)
  54.                 x_load.append(x)
  55.                 y_load.append(labels.index(cat))  # directory name as label
  56.             except Exception as e:
  57.                 print(e)
  58.                 continue
  59.     return np.array(x_load), np.array(y_load)
  61. def prepress_labels(labels):
  62.     """
  63.     # one-hot编码 把类别id转换为表示当前类别的向量，比如0 1 2 =》 [[1 0 0] [0 1 0] [0 0 1]]
  64.     """
  65.     from tensorflow.keras import utils
  66.     labels = utils.to_categorical(labels)
  67.     return labels
  69. class AccuracyLogger(keras.callbacks.Callback):
  70.     """
  71.     AccuracyLogger 类, 用来记录训练进度
  72.     Attributes:
  73.     callback: function 回调函数
  74.     epoch: int 回合
  75.     batch: int 批次
  76.     """
  78.     def __init__(self, callback=None):
  79.         self.callback = callback
  80.         self.epoch = 0
  81.         self.batch = 0
  83.     def set_model(self, model):
  84.         """
  85.         set_model 函数, 在训练之前由父模型调用，告诉回调函数是哪个模型在调用它
  86.         Keyword arguments::
  87.         model: model 训练模型
  88.         """
  89.         self.model = model  #
  91.     def on_epoch_begin(self, epoch, logs={}):
  92.         """
  93.         on_epoch_begin 函数, epoch开始的时候自动调用
  94.         Keyword arguments::
  95.         epoch: int 回合数
  96.         """
  97.         self.epoch = epoch
  99.     def on_batch_end(self, batch, logs={}):
  100.         """
  101.         on_batch_end 函数, 在batch训练结束后自动调用
  102.         Keyword arguments::
  103.         batch: int 批次数
  104.         """
  105.         self.batch = batch
  106.         print('epoch:%d batch:%d' % (self.epoch, self.batch))
  107.         # print(logs)
  108.         if self.callback is not None:
  109.             self.callback(self.epoch, self.batch, logs)
  112. class ImageClassifier:
  113.     """
  114.     ImageClassifier 类, 用来训练分类器
  115.     Attributes:
  116.     data_root: str 数据根目录
  117.     FC_NUMS: int 回合
  118.     TRAIN_LAYERS: int 批次
  119.     model: model 模型
  120.     busy: bool 是否正在执行操作
  121.     ns: str 单次训练命名空间
  122.     """
  124.     def __init__(self, model_path=None, data_root=os.path.expanduser('~')+'/Lepi_Data/ros/transfer_learning'):
  125.         self.data_root = data_root
  126.         self.FC_NUMS = 64
  127.         self.TRAIN_LAYERS = 2
  128.         self.IMAGE_SIZE = 112
  129.         self.model = None
  130.         self.busy = False
  131.         self.ns = None
  133.     def load_model(self):
  134.         """
  135.         on_batch_end 函数, 在batch训练结束后自动调用
  136.         Keyword arguments::
  137.         batch: int 批次数
  138.         """
  139.         self.busy = True
  140.         data_dir = os.path.join(self.data_root, self.ns)
  141.         path = os.path.join(data_dir, 'model.h5')
  142.         try:
  143.             if self.model is None:
  144.                 self.model = load_model(path)  # 加载训练模型
  145.             else:
  146.                 K.clear_session()
  147.                 self.model = load_model(path)
  148.         except Exception as e:
  149.             print(e)
  150.         finally:
  151.             self.busy = False
  153.     def preprocess_input(self, cv_img):
  154.         cv_img = cv2.resize(
  155.             cv_img, (self.IMAGE_SIZE, self.IMAGE_SIZE)).astype('float32')
  156.         # cv_img = cv2.normalize(cv_img, None, 0, 1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32FC3)
  157.         # cv_img = cv_img.reshape(-1, self.IMAGE_SIZE, self.IMAGE_SIZE, 3)
  158.         input_data = np.expand_dims(cv_img, axis=0)
  159.         return input_data
  161.     def get_base_model(self, name='vgg16'):
  162.         print('using model : ', name)
  163.         from tensorflow.keras.applications.vgg16 import VGG16
  164.         # 采用VGG16为基本模型，include_top为False，表示FC层是可自定义的，抛弃模型中的FC层；该模型会在~/.keras/models下载基本模型
  165.         base_model = VGG16(input_shape=(
  166.             self.IMAGE_SIZE, self.IMAGE_SIZE, 3), include_top=False, weights='imagenet')
  167.         # self.preprocess_input = preprocess_input
  168.         return base_model
  170.     def download_model(self, name='vgg16'):
  171.         """
  172.         下载模型
  173.         """
  174.         self.get_base_model(name)
  176.     def train(self, data_dir, epochs=3, callback=None, model_name='vgg16'):
  177.         """
  178.         train 函数, 训练模型
  179.         Keyword arguments::
  180.         epochs: int 训练次数
  181.         """
  182.         if self.busy:
  183.             return 1
  184.         self.busy = True
  185.         label_names = get_labels(data_dir)
  186.         self.NUM_CLASSES = len(label_names)
  187.         base_model = self.get_base_model(model_name)
  188.         x_data, y_label = load_img_from_dir(data_dir, target_size=(
  189.             self.IMAGE_SIZE, self.IMAGE_SIZE), max_num=30)
  190.         for i in range(x_data.shape[0]):
  191.             x_data[i] = self.preprocess_input(x_data[i])
  192.         print(x_data.shape)
  193.         print(x_data[0].shape)
  194.         x_data = x_data.reshape(
  195.             x_data.shape[0], self.IMAGE_SIZE, self.IMAGE_SIZE, 3)
  196.         y_label_one_hot = prepress_labels(y_label)
  197.         # 验证应该使用从未见过的图片
  198.         train_x, test_x, train_y, test_y = train_test_split(x_data, y_label_one_hot, random_state=0,
  199.                                                             test_size=0.3)
  200.         # 自定义FC层以基本模型的输入为卷积层的最后一层
  201.         x = base_model.output
  202.         x = GlobalAveragePooling2D()(x)
  203.         x = Dense(self.FC_NUMS, activation='relu')(x)
  204.         prediction = Dense(self.NUM_CLASSES, activation='softmax')(x)
  206.         # 构造完新的FC层，加入custom层
  207.         model = Model(inputs=base_model.input, outputs=prediction)
  209.         # 获取模型的层数
  210.         print("layer nums:", len(model.layers))
  212.         # 除了FC层，靠近FC层的一部分卷积层可参与参数训练，
  213.         # 一般来说，模型结构已经标明一个卷积块包含的层数，
  214.         # 在这里我们选择TRAIN_LAYERS为3，表示最后一个卷积块和FC层要参与参数训练
  215.         for layer in model.layers:
  216.             layer.trainable = False
  217.         for layer in model.layers[-self.TRAIN_LAYERS:]:
  218.             layer.trainable = True
  219.         for layer in model.layers:
  220.             print("layer.trainable:", layer.trainable)
  222.         # 预编译模型
  223.         model.compile(optimizer=SGD(lr=0.0001, momentum=0.9),
  224.                       loss='categorical_crossentropy', metrics=['accuracy'])
  225.         model.summary()
  226.         model.fit(train_x, train_y,
  227.                   validation_data=(test_x, test_y),
  228.                   # model.fit(x_data, y_label_one_hot,
  229.                   #         validation_split=0.4,
  230.                   callbacks=[AccuracyLogger(callback)],
  231.                   epochs=epochs, batch_size=4,
  232.                   # steps_per_epoch=1,validation_steps =1 ,
  233.                   verbose=1, shuffle=True)
  234.         self.model = model
  235.         model.save(os.path.join(data_dir, 'model.h5'))
  236.         self.label_names = label_names
  237.         self.dump_label_name(label_names)
  238.         # self.convert_tflite()
  239.         # self.session = K.get_session()
  240.         # self.graph = tf.get_default_graph()
  241.         self.busy = False
  243.     def dump_label_name(self, dirs):
  244.         """
  245.         保存训练标签
  246.         """
  247.         with open(os.path.join(self.data_root, self.ns, "labelmap.txt"), "w") as f:
  248.             # pickle.dump(dirs, f, protocol=2)
  249.             to_write = []
  250.             for item in dirs:
  251.                 to_write.append(item+'\n')
  252.             f.writelines(to_write)
  254.     def load_label_name(self):
  255.         """
  256.         加载训练标签
  257.         """
  258.         import numpy as np
  259.         with open(os.path.join(self.data_root, self.ns, "labelmap.txt"), "r") as f:
  260.             # label_names = np.array(pickle.load(f))
  261.             label_names = [line.strip() for line in f.readlines()]
  262.         self.label_names = label_names
  263.         return label_names
  265.     def evaluate(self, data_dir):
  266.         """
  267.         模型评估函数
  268.         """
  269.         x_data, y_label = load_img_from_dir(data_dir, target_size=(
  270.             self.IMAGE_SIZE, self.IMAGE_SIZE), max_num=60)
  271.         for i in range(x_data.shape[0]):
  272.             x_data[i] = self.preprocess_input(x_data[i])
  273.         x_data = x_data.reshape(
  274.             x_data.shape[0], self.IMAGE_SIZE, self.IMAGE_SIZE, 3)
  275.         y_label_one_hot = prepress_labels(y_label)
  276.         # 验证应该使用从未见过的图片
  277.         train_x, test_x, train_y, test_y = train_test_split(x_data, y_label_one_hot, random_state=0,
  278.                                                             test_size=0.5)
  279.         # 开始评估模型效果 # verbose=0为不输出日志信息
  280.         score = model.evaluate(test_x, test_y, verbose=1, steps=1)
  281.         print('Test loss:', score[0])
  282.         print('Test accuracy:', score[1])  # 准确度
  284.     def predict(self, cv_img=None, path=None):
  285.         """
  286.         predict 函数, 执行预测
  287.         Keyword arguments::
  288.         cv_img: image 输入opencv图像
  289.         path: str 图片路径
  290.         """
  291.         if path is not None:
  292.             cv_img = cv2.imread(path)
  293.         # label_names = get_labels(data_dir)
  294.         # rs_img_f32 = cv2.resize(cv_img, (self.IMAGE_SIZE, self.IMAGE_SIZE)).astype('float32')
  295.         if cv_img is None:
  296.             print('no cv image provided')
  297.             return None
  298.         input_data = self.preprocess_input(cv_img)
  299.         if self.model is not None:
  300.             result = self.model.predict(input_data, steps=1)
  301.             # print(label_names)
  302.             print("result:", result)
  303.             return self.label_names, result[0]
  304.         else:
  305.             print('your model is not ready')
  306.             return None
  309. if __name__ == '__main__':
  310.     """
  311.     训练测试
  312.     """
  313.     IC = ImageClassifier()
  314.     IC.ns = '分类测试'
  315.     epoch_total = 3
  317.     def pub_training_logs(epoch, batch, logs):
  318.         # logs {'loss': 0.33773628, 'accuracy': 0.71428573, 'batch': 6, 'size': 4}
  319.         print(logs)
  320.         msg = '第%d/%d轮, 批次: %d, 损失: %.2f, 准确率: %.2f' % (
  321.             epoch+1, epoch_total, batch, logs['loss'], logs['accuracy'])
  322.         print(msg)
  323.     IC.train(os.path.join(IC.data_root, '分类测试'), epoch_total,
  324.              callback=pub_training_logs, model_name='vgg16')