BIDI_ML_INFO_EXTRACTION/BiddingKG/dl/product/train.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Author  : bidikeji
# @Time    : 2021/12/14 11:08

#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Author  : bidikeji
# @Time    : 2021/1/13 0013 10:12
import os
import re
print('准备导入tf')
import tensorflow as tf
print('准备导入np')
import numpy as np
print('准备导入Product_model')
from BiddingKG.dl.product.product_model import Product_Model
print('准备导入max_len')
from BiddingKG.dl.product.data_tfrecord import max_len
# from BiddingKG.dl.common.Utils import viterbi_decode
print('准备设置CUDA环境')
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
# max_len = 500

def read_tfRecord(sess, file_tfRecord):
    queue = tf.train.string_input_producer([file_tfRecord])
    reader = tf.TFRecordReader()
    filename_, serialized_example = reader.read(queue)
    features = tf.parse_single_example(
        serialized_example,
        features = {
            "word_ids": tf.FixedLenFeature([max_len], tf.int64),
            "text_len":tf.FixedLenFeature([1], tf.int64),
            "tag_ids": tf.FixedLenFeature([max_len], tf.int64)
        }
    )
    word_ids = tf.cast(features['word_ids'], tf.int64)
    tag_ids = tf.cast(features['tag_ids'], tf.int64)
    text_len = tf.cast(features['text_len'], tf.int64)
    return text_len, word_ids , tag_ids

def get_batch_record(sess,filename, batch_size):
    text_len, word_ids, tag_ids = read_tfRecord(sess, filename)
    text_len, word_ids, tag_ids = tf.train.shuffle_batch([text_len, word_ids , tag_ids],
                                                     batch_size=batch_size,
                                                     capacity=200+batch_size*3,
                                                     min_after_dequeue=1,
                                                     num_threads=5)
    text_len = tf.squeeze(text_len, squeeze_dims=1)
    return text_len, word_ids , tag_ids

def total_sample(file_name):
    '''
    返回样本总数
    :param file_name: tfrecord 文件名
    :return:
    '''
    sample_num = 0
    for record in tf.python_io.tf_record_iterator(file_name):
        sample_num += 1
    return sample_num

if __name__ == "__main__":
    print('进入main ')
    filename = os.path.dirname(__file__)+'/data/train_data/maxlen_500_addunk_product_reason.tfrecode'
    filename_dev = os.path.dirname(__file__)+'/data/test_data/maxlen_500_addunk_product_reason.tfrecode'
    assert os.path.exists(filename)
    assert os.path.exists(filename_dev)
    print('确保文件存在')
    batch_size = 100
    # id_to_tag = {0: 'O', 1: 'B', 2: 'I', 3: 'E'}
    tag2index = {'S': 0, 'B-pro': 1, 'I-pro': 2, 'E-pro': 3, 'B-rea': 4, 'I-rea': 5, 'E-rea': 6}
    id_to_tag = {v:k for k,v in tag2index.items()}

    def get_ner(id_tag):
        ner = set()
        for it in re.finditer('12*3', id_tag):
            ner.add((it.start(), it.end(), 'pro'))
        for it in re.finditer('45*6', id_tag):
            ner.add((it.start(), it.end(), 'rea'))
        return ner

    print('准备构建模型')
    model = Product_Model()
    # saver = tf.train.Saver()
    print('模型构建完毕')
    with tf.Session() as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        print('参数初始化')
        text_len, word_ids, tag_ids = get_batch_record(sess, filename, batch_size=batch_size)
        print('get_batch_record')
        text_len_dev, word_ids_dev, tag_ids_dev = get_batch_record(sess, filename_dev, batch_size=batch_size)
        print('get_batch_record_dev')
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        print('total_sample(filename)', total_sample(filename))

        total_num = total_sample(filename)
        batch_num = total_num//batch_size
        batch_num_dev = total_sample(filename_dev)//batch_size
        num = 0
        l = []
        max_f1 = 0

        # model.saver.restore(sess, os.path.join(os.path.dirname(__file__)+'/model','ner_epoch10_f10.6875_loss1.5230.ckpt'))
        # print('模型加载成功')

        for epoch in range(50):
            for batch in range(batch_num):
                text_len_, word_ids_, tag_ids_ = sess.run([text_len, word_ids, tag_ids])
                # print(text_len_.shape, word_ids_.shape, tag_ids_.shape)
                loss_, train_op_, globel_step_ = sess.run([model.loss, model.train_op, model.globel_step],
                                                                 feed_dict={model.char_inputs:word_ids_,
                                                        # model.lengths:text_len_,
                                                        model.dropout:0.5,
                                                        model.targets:tag_ids_})


                if batch % 100==0:
                    print('loss_:', loss_, '\tglobel_step_:',globel_step_)
                num += text_len_.shape[0]
            print('训练总数：%d, 样本总数：%d'%(num, total_num))

            results = []
            trans = model.trans.eval()
            Precision = []
            Recall = []
            F1 = []
            loss = []
            pred_num = 0
            gold_num = 0
            equal_num = 0
            for batch in range(batch_num_dev):
                text_len_dev_, word_ids_dev_, tag_ids_dev_ = sess.run([text_len_dev, word_ids_dev, tag_ids_dev])
                lengths, scores, batch_loss = sess.run([model.lengths, model.logits, model.loss], feed_dict={model.char_inputs:word_ids_dev_,
                                                        # model.lengths:text_len_dev_,
                                                        model.dropout:0.5,
                                                        model.targets:tag_ids_dev_})
                loss.append(batch_loss)
                batch_paths = model.decode(scores, lengths, trans)
                for i in range(len(tag_ids_dev_)):
                    result = []
                    gold = [str(x) for x in tag_ids_dev_[i][:lengths[i]]]
                    pred = [str(x) for x in batch_paths[i][:lengths[i]]]
                    gold_ner = get_ner("".join(gold))
                    pred_ner = get_ner("".join(pred))
                    # print('标签实体：',gold_ner)
                    # print('预测实体：',pred_ner)
                    pred_num += len(pred_ner)
                    gold_num += len(gold_ner)
                    equal_num += len(gold_ner & pred_ner)
            precision = equal_num / (pred_num + 1e-10)
            recall = equal_num / (gold_num + 1e-10)
            f1 = 2 * (precision * recall) / (precision + recall + 1e-10)
            val_loss = np.mean(loss)
            print('epoch: %d, f1:%.4f, acc:%.4f, recall:%.4f, val_loss:%.4f'%(epoch, f1, precision, recall, val_loss))
            if f1>max_f1:
                max_f1 = f1
                model.saver.save(sess, os.path.join(os.path.dirname(__file__)+'/model', "ner_epoch%d_f1%.4f_loss%.4f.ckpt"%(epoch,f1, val_loss)))
                print('save model, max_f1:%.4f' %f1)

        coord.request_stop()
        coord.join(threads)