def __init__(self): self.config = TCNNConfig() self.categories, self.cat_to_id = read_category() self.words, self.word_to_id = read_vocab(vocab_dir) self.config.vocab_size = len(self.words) self.model = TextCNN(self.config) self.session = tf.Session() self.session.run(tf.global_variables_initializer()) saver = tf.train.Saver() saver.restore(sess=self.session, save_path=save_path)
# if len(sys.argv) != 2 or sys.argv[1] not in ['train', 'test']: # raise ValueError("""usage: python run_rnn.py [train / test]""") d = 'data/cnews/' train_dir = d + 'sen_class.train' test_dir = d + 'sen_class.test' val_dir = d + 'sen_class.val' vocab_dir = d + 'sen_class.vocab' print('Configuring RNN model...') config = TRNNConfig() labels = build_vocab(train_dir, test_dir, val_dir, vocab_dir, config.vocab_size) config.num_classes = len(labels) categories, cat_to_id = read_category(labels) words, word_to_id = read_vocab(vocab_dir) config.vocab_size = len(words) model = TextRNN(config) print('labels: {}, vocabulary size: {}'.format(config.num_classes, len(words))) with open(map_path, "wb") as f: pickle.dump( [word_to_id, cat_to_id, config.seq_length, config.num_classes], f) train() exit() if sys.argv[1] == 'train': train() elif sys.argv[1] == 'test': test()