main.py

import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
import model as ml
import data
import numpy as np
import os
import sys
from configs import DEFINES
DATA_OUT_PATH = './data_out/'
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--train', action='store_true', default=False, help='for training')
parser.add_argument('--test', action='store_true', default=False, help='for testing')
parser.add_argument('--train_steps', type=int, default=100000, help='for training')
parser.add_argument('--batch_size', type=int, default=4, help='for training')
args = parser.parse_args()

def main(self):
    data_out_path = os.path.join(os.getcwd(), DATA_OUT_PATH)
    os.makedirs(data_out_path, exist_ok=True)
    # 데이터를 통한 사전 구성 한다.
    char2idx, idx2char, vocabulary_length = data.load_vocabulary()

    # 현재 경로'./'에 현재 경로 하부에 
    # 체크 포인트를 저장한 디렉토리를 설정한다.
    check_point_path = os.path.join(os.getcwd(), DEFINES.check_point_path)
    # 디렉토리를 만드는 함수이며 두번째 인자 exist_ok가 
    # True이면 디렉토리가 이미 존재해도 OSError가 
    # 발생하지 않는다.
    # exist_ok가 False이면 이미 존재하면 
    # OSError가 발생한다.
    os.makedirs(check_point_path, exist_ok=True)
    # ----------------------------------
    tf.keras.backend.clear_session()
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    # ---------------------------------
    # 에스티메이터 구성한다.
    classifier = tf.estimator.Estimator(
        model_fn=ml.Model,  # 모델 등록한다.
        model_dir=DEFINES.check_point_path,  # 체크포인트 위치 등록한다.
        params={  # 모델 쪽으로 파라메터 전달한다.
            'embedding_size': DEFINES.embedding_size,
            'model_hidden_size': DEFINES.model_hidden_size,  # 가중치 크기 설정한다.
            'ffn_hidden_size': DEFINES.ffn_hidden_size,
            'attention_head_size': DEFINES.attention_head_size,
            'learning_rate': DEFINES.learning_rate,  # 학습율 설정한다.
            'vocabulary_length': vocabulary_length,  # 딕셔너리 크기를 설정한다.
            'embedding_size': DEFINES.embedding_size,  # 임베딩 크기를 설정한다.
            'layer_size': DEFINES.layer_size,
            'max_sequence_length': DEFINES.max_sequence_length,
            'xavier_initializer': DEFINES.xavier_initializer
        })

    if args.train:
        # 훈련 데이터와 테스트 데이터를 가져온다.
        train_input, train_label, eval_input, eval_label = data.load_data()

        # 훈련셋 인코딩 만드는 부분이다.
        train_input_enc, train_input_enc_length = data.enc_processing(train_input, char2idx)
        # 훈련셋 디코딩 입력 부분 만드는 부분이다.
        train_output_dec, train_output_dec_length = data.dec_output_processing(train_label, char2idx)
        # 훈련셋 디코딩 출력 부분 만드는 부분이다.
        train_target_dec = data.dec_target_processing(train_label, char2idx)

        # 평가셋 인코딩 만드는 부분이다.
        eval_input_enc, eval_input_enc_length = data.enc_processing(eval_input, char2idx)
        # 평가셋 인코딩 만드는 부분이다.
        eval_output_dec, eval_output_dec_length = data.dec_output_processing(eval_label, char2idx)
        # 평가셋 인코딩 만드는 부분이다.
        eval_target_dec = data.dec_target_processing(eval_label, char2idx)
        # 학습 실행
        classifier.train(input_fn=lambda: data.train_input_fn(
            train_input_enc, train_output_dec, train_target_dec, args.batch_size), steps=args.train_steps)
    
    # test 부분
    if args.test:
        q_sent = input('Review sentence > ').strip()
        predic_input_enc, predic_input_enc_length = data.enc_processing([q_sent], char2idx)
        # 학습 과정이 아니므로 디코딩 입력은
        # 존재하지 않는다.(구조를 맞추기 위해 넣는다.)
        predic_output_dec, predic_output_decLength = data.dec_output_processing([""], char2idx)
        # 학습 과정이 아니므로 디코딩 출력 부분도
        # 존재하지 않는다.(구조를 맞추기 위해 넣는다.)
        predic_target_dec = data.dec_target_processing([""], char2idx)

        
        for i in range(DEFINES.max_sequence_length):
            if i > 0:
                predic_output_dec, predic_output_decLength = data.dec_output_processing([answer], char2idx)
                predic_target_dec = data.dec_target_processing([answer], char2idx)
            # 예측을 하는 부분이다.
            predictions = classifier.predict(
                input_fn=lambda: data.eval_input_fn(predic_input_enc, predic_output_dec, predic_target_dec, 1))

            answer, finished = data.pred_next_string(predictions, idx2char)
            if finished:
                break
        print("summary: ", answer)
        '''
        predictions = classifier.predict(
            input_fn=lambda: data.eval_input_fn(predic_input_enc, predic_output_dec, predic_target_dec, 1))

        answer, finished = data.pred_next_string(predictions, idx2char)
        # 예측한 값을 인지 할 수 있도록
        # 텍스트로 변경하는 부분이다.
        print("answer: ", answer)
        '''


if __name__ == '__main__':
    #tf.logging.set_verbosity(tf.logging.INFO)
    tf.app.run(main)