BIDI_ML_INFO_EXTRACTION/BiddingKG/dl/table_head/train.py

import sys
import os
sys.path.append(os.path.abspath("../../.."))
os.environ['KERAS_BACKEND'] = 'tensorflow'
from keras.metrics import categorical_accuracy
from BiddingKG.dl.table_head.metrics import precision, recall, f1
from keras import optimizers, Model
from BiddingKG.dl.table_head.models.model import get_model
from BiddingKG.dl.table_head.loss import focal_loss
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
from BiddingKG.dl.table_head.pre_process import get_data_from_file, get_data_from_sql, my_data_loader
from keras import backend as K


input_shape = (6, 10, 60)
output_shape = (1,)
batch_size = 32
epochs = 1000
pretrained_path = "checkpoints/best.hdf5"
checkpoint_path = "checkpoints/"
PRETRAINED = True
CHECKPOINT = False


def train():
    # GPU available
    print("gpus", K.tensorflow_backend._get_available_gpus())

    # Data
    data_x, data_y = get_data_from_file('txt')
    # data_x = data_x[:60000]
    # data_y = data_y[:60000]
    print("finish read data", len(data_x))

    # Split -> Train, Test
    split_size = int(len(data_x)*0.1)
    test_x, test_y = data_x[:split_size], data_y[:split_size]
    train_x, train_y = data_x[split_size:], data_y[split_size:]

    # Data Loader
    train_data_loader = my_data_loader(train_x, train_y, batch_size=batch_size)
    test_data_loader = my_data_loader(test_x, test_y, batch_size=batch_size)

    # Model
    model = get_model(input_shape, output_shape)
    if PRETRAINED:
        model.load_weights(pretrained_path)
        print("read pretrained model", pretrained_path)
    else:
        print("no pretrained")
    if CHECKPOINT:
        model.load_weights(checkpoint_path)
        print("read checkpoint model", checkpoint_path)
    else:
        print("no checkpoint")

    filepath = 'e{epoch:02d}-f1{val_f1:.2f}'
    checkpoint = ModelCheckpoint(checkpoint_path+filepath+".hdf5", monitor='val_f1',
                                 verbose=1, save_best_only=True, mode='max')

    model.compile(optimizer=optimizers.Adam(lr=0.005), loss=focal_loss(),
    # model.compile(optimizer=optimizers.Adam(lr=0.005), loss='binary_crossentropy',
                  metrics=['acc',
                           precision, recall, f1])

    rlu = ReduceLROnPlateau(monitor='val_f1', factor=0.1, patience=5,
                            verbose=1, mode='max', cooldown=0, min_lr=0)

    model.fit_generator(train_data_loader,
                        steps_per_epoch=max(1, len(train_x) // batch_size),
                        callbacks=[checkpoint, rlu],
                        validation_data=test_data_loader,
                        validation_steps=max(1, len(test_x) // batch_size),
                        epochs=epochs)

    # model.fit(x=[train_x[0], train_x[1], train_x[2]], y=train_y,
    #           validation_data=([test_x[0], test_x[1], test_x[2]], test_y),
    #           epochs=epochs, batch_size=256, shuffle=True,
    #           callbacks=[checkpoint, rlu])

    return model, test_x


def print_layer_output(model, data):
    middle_layer = Model(inputs=model.inputs,
                         outputs=model.get_layer('input_2').output)

    middle_layer_output = middle_layer.predict([data[0], data[1]])
    print(middle_layer_output)
    return


if __name__ == '__main__':
    model, data = train()

    # place_list = get_place_list()
    # _str1 = '中国电信'
    # _str2 = '分公司'
    # _list = []
    # for place in place_list:
    #     _list.append(_str1 + place + _str2 + "\n")
    # # print(_list)
    # with open("电信分公司.txt", "w") as f:
    #     f.writelines(_list)