[人工智能]Tensorflow2.0实现断点续训

参考: https://www.bilibili.com/video/BV16A41157LW?p=17
视频及课件来源 北京大学 曹建
使用的识别图片

获取训练数据集

def get_mnist_data():
    # 参考: https://www.codenong.com/53310656/
    # 获取数据 return  (x_train, y_train), (x_test, y_test)
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    # 数据归一化 0-255之间的灰度值变成 变成0 或者 1
    # 把输入特征的数值变小 更利于神经网络的吸收
    x_train, x_test = (x_train > 128).astype(int), (x_test > 128).astype(int)
    return (x_train, y_train), (x_test, y_test)

加载训练模型

def local_load_weights(checkpoint_save_path):
    local_weights_model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    local_weights_model.compile(optimizer='adam',
                                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
                                metrics=['sparse_categorical_accuracy'])
    if os.path.exists(checkpoint_save_path + '.index'):
        tf.print('-------------load the model-----------------')
        local_weights_model.load_weights(checkpoint_save_path)
    return local_weights_model

断点续练

def weights_train(local_weights_model, checkpoint_save_path):
    # 回调函数 用于保存模型
    cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,  # 文件存储路径
                                                     save_weights_only=True,  # 是否只保留模型参数
                                                     save_best_only=True)  # 是否只保留最优结果
    (x_train, y_train), (x_test, y_test) = get_mnist_data()
    history = local_weights_model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_test, y_test),
                                      validation_freq=1,
                                      callbacks=[cp_callback])  # fit中假如回调函数

还有一种方法
加载模型

def local_load_model(model_path):
    if os.path.exists(model_path + '/saved_model.pb'):
        tf.print('-------------load the model-----------------')
        local_model = tf.keras.models.load_model(model_path)
    else:
        local_model = tf.keras.models.Sequential([
            tf.keras.layers.Flatten(),  # 拉直层  将数据拉直成1维
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        local_model.compile(optimizer='adam',
                            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
                            metrics=['sparse_categorical_accuracy'])
    return local_model

训练

 (x_train, y_train), (x_test, y_test) = get_mnist_data()
    history = local_model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_test, y_test),
                              validation_freq=1)

整个文件如下

# -*- coding: utf-8 -*-

import tensorflow as tf
import os
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt


def get_mnist_data():
    # 参考: https://www.codenong.com/53310656/
    # 获取数据 return  (x_train, y_train), (x_test, y_test)
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    # 数据归一化 0-255之间的灰度值变成 变成0 或者 1
    # 把输入特征的数值变小 更利于神经网络的吸收
    x_train, x_test = (x_train > 128).astype(int), (x_test > 128).astype(int)
    return (x_train, y_train), (x_test, y_test)


def get_local_image():
    img = Image.open('./data/image/2.jpg')  # 载入自己的图片
    img = img.resize((28, 28))  # 设置图片大小
    gray_img = img.convert('L')
    mun_img = np.array(gray_img)
    # convert.show()
    mun_img = (mun_img > 64).astype(int)
    x_test__reshape = mun_img.reshape(1, 28, 28)
    return x_test__reshape


def local_load_weights(checkpoint_save_path):
    local_weights_model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    local_weights_model.compile(optimizer='adam',
                                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
                                metrics=['sparse_categorical_accuracy'])
    if os.path.exists(checkpoint_save_path + '.index'):
        tf.print('-------------load the model-----------------')
        local_weights_model.load_weights(checkpoint_save_path)
        tf.print("加载文件参数")
    return local_weights_model


def local_load_model(model_path):
    if os.path.exists(model_path + '/saved_model.pb'):
        tf.print('-------------load the model-----------------')
        local_model = tf.keras.models.load_model(model_path)
    else:
        local_model = tf.keras.models.Sequential([
            tf.keras.layers.Flatten(),  # 拉直层  将数据拉直成1维
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        local_model.compile(optimizer='adam',
                            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
                            metrics=['sparse_categorical_accuracy'])
    return local_model


# 断点续练
def weights_train(local_weights_model, checkpoint_save_path):
    # 回调函数 用于保存模型
    cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,  # 文件存储路径
                                                     save_weights_only=True,  # 是否只保留模型参数
                                                     save_best_only=True)  # 是否只保留最优结果
    (x_train, y_train), (x_test, y_test) = get_mnist_data()
    history = local_weights_model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_test, y_test),
                                      validation_freq=1,
                                      callbacks=[cp_callback])  # fit中假如回调函数
    local_weights_model.summary()
    show_train_line(history)


def show_train_line(history):
    # 显示训练集和验证集的acc和loss曲线
    acc = history.history['sparse_categorical_accuracy']
    val_acc = history.history['val_sparse_categorical_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    plt.subplot(1, 2, 1)
    plt.plot(acc, label='Training Accuracy')
    plt.plot(val_acc, label='Validation Accuracy')
    plt.title('Training and Validation Accuracy')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(loss, label='Training Loss')
    plt.plot(val_loss, label='Validation Loss')
    plt.title('Training and Validation Loss')
    plt.legend()
    plt.show()


def forecast_demo(model):
    # 预测
    local_image = get_local_image()
    loaded_evaluate = model.predict(local_image)
    print(loaded_evaluate)
    prediction = np.argmax(loaded_evaluate, axis=1)  # 找出最大值
    print('预测结果：', prediction)


def demo_1():
    checkpoint_path = "./data/model/checkpoint/mnist.ckpt"
    local_weights_model = local_load_weights(checkpoint_path)
    # 继续训练
    weights_train(local_weights_model, checkpoint_path)
    forecast_demo(local_weights_model)


def demo_2():
    model_path = "./data/model/breakpoint"
    local_model = local_load_model(model_path)
    (x_train, y_train), (x_test, y_test) = get_mnist_data()
    history = local_model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_test, y_test),
                              validation_freq=1)
    show_train_line(history)
    local_model.summary()
    #   保存模型
    local_model.save(model_path, save_format="tf")
    forecast_demo(local_model)


if __name__ == '__main__':
    demo_1()
    # demo_2()