LSTM.py

# coding:utf-8
"""
Created on Tue Apr 23 14:37:41 2019
@author: jiali zhang
bearing fault diagnosis by LSTM
"""

from keras.layers import Dense, Activation, Flatten, LSTM
from keras.models import Sequential
from keras.utils import plot_model
# from keras.regularizers import l2
import preprocess
from keras.callbacks import TensorBoard
import numpy as np
import time

for i in range(0,4):
    # 训练参数
    batch_size = 128
    epochs = 12
    num_classes = 10
    length = 2048
    BatchNorm = True  # 是否批量归一化
    number = 1000  # 每类样本的数量
    normal = True  # 是否标准化
    rate = [0.6, 0.2, 0.2]  # 测试集验证集划分比例
    date = time.strftime("%Y%m%d", time.localtime())
    mark = time.strftime("%Y%m%d_%H%M", time.localtime())

    path = r'data\0HP'
    x_train, y_train, x_valid, y_valid, x_test, y_test = preprocess.prepro(d_path=path, length=length,
                                                                           number=number,
                                                                           normal=normal,
                                                                           rate=rate,
                                                                           enc=True, enc_step=28)

    x_train, x_valid, x_test = x_train[:, :, np.newaxis], x_valid[:, :, np.newaxis], x_test[:, :, np.newaxis]

    input_shape = x_train.shape[1:]

    print('训练样本维度:', x_train.shape)
    print(x_train.shape[0], '训练样本个数')
    print('验证样本的维度', x_valid.shape)
    print(x_valid.shape[0], '验证样本个数')
    print('测试样本的维度', x_test.shape)
    print(x_test.shape[0], '测试样本个数')

    model_name = "lstm_diagnosis-{}".format(mark)

    # 实例化一个Sequential
    model = Sequential()

    model.add(LSTM(32, activation='tanh', recurrent_activation='hard_sigmoid', kernel_initializer='glorot_uniform',
                   recurrent_initializer='orthogonal', bias_initializer='zeros', return_sequences=True))

    model.add(Flatten())

    # 增加输出层，共num_classes个单元，激活函数为softmax
    # model.add(Dense(units=num_classes, activation='softmax', kernel_regularizer=l2(1e-4)))
    model.add(Dense(units=num_classes, activation='softmax'))

    # 编译模型 评价函数和损失函数相似，不过评价函数的结果不会用于训练过程中
    model.compile(optimizer='Adam', loss='categorical_crossentropy',
                  metrics=['accuracy'])

    # TensorBoard调用查看一下训练情况
    tb_cb = TensorBoard(log_dir='logs\{}_logs\{}'.format(date, model_name))

    # 开始模型训练
    model.fit(x=x_train, y=y_train, batch_size=batch_size, epochs=epochs,
              verbose=1, validation_data=(x_valid, y_valid), shuffle=True,
              callbacks=[tb_cb])

    # 评估模型
    score = model.evaluate(x=x_test, y=y_test, verbose=0)
    print("测试集上的损失率：", score[0])
    print("测试集上的准确率：", score[1])
    plot_model(model=model, to_file='lstm-diagnosis.png', show_shapes=True)