nvdm.py 文件源码

def __init__(self, sess, reader, dataset="ptb",
               decay_rate=0.96, decay_step=10000, embed_dim=500,
               h_dim=50, learning_rate=0.001, max_iter=450000,
               checkpoint_dir="checkpoint"):
    """Initialize Neural Varational Document Model.

    params:
      sess: TensorFlow Session object.
      reader: TextReader object for training and test.
      dataset: The name of dataset to use.
      h_dim: The dimension of document representations (h). [50, 200]
    """
    self.sess = sess
    self.reader = reader

    self.h_dim = h_dim
    self.embed_dim = embed_dim

    self.max_iter = max_iter
    self.decay_rate = decay_rate
    self.decay_step = decay_step
    self.checkpoint_dir = checkpoint_dir
    self.step = tf.Variable(0, trainable=False)  
    self.lr = tf.train.exponential_decay(
        learning_rate, self.step, 10000, decay_rate, staircase=True, name="lr")

    _ = tf.scalar_summary("learning rate", self.lr)

    self.dataset = dataset
    self._attrs = ["h_dim", "embed_dim", "max_iter", "dataset",
                   "learning_rate", "decay_rate", "decay_step"]

    self.build_model()