model.py 文件源码

def __init__(self,params):
        config = tf.ConfigProto(allow_soft_placement=True)
        self.sess = tf.Session(config = config)
        K.set_session(self.sess)
        # Pull out all of the parameters
        self.batch_size = params['batch_size']
        self.seq_len = params['seq_len']
        self.vocab_size = params['vocab_size']
        self.embed_size = params['embed_size']
        self.hidden_dim = params['hidden_dim']
        self.num_layers = params['num_layers']
        with tf.device('/gpu:0'):
            # Set up the input placeholder
            self.input_seq = tf.placeholder(tf.float32, shape=[None, self.seq_len])
            # Build the RNN
            self.rnn = Embedding(self.vocab_size + 1, self.embed_size, input_length=self.seq_len)(self.input_seq)
        with tf.device('/gpu:1'):
            for l in range(self.num_layers):
                self.rnn = LSTM(output_dim=self.hidden_dim, return_sequences=True, name='rnn_1')(self.rnn)
            rnn_output = tf.unpack(self.rnn, axis=1)
            self.w_proj = tf.Variable(tf.zeros([self.vocab_size, self.hidden_dim]))
            self.b_proj = tf.Variable(tf.zeros([self.vocab_size]))
            self.output_seq = tf.placeholder(tf.int64, shape=([None, self.seq_len]))
            losses = []
            outputs = []
            for t in range(self.seq_len):
                rnn_t = rnn_output[t]
                y_t = tf.reshape(self.output_seq[:, t],[-1,1])
                step_loss = tf.nn.sampled_softmax_loss(weights=self.w_proj, biases=self.b_proj, inputs=rnn_t,
                                                       labels=y_t, num_sampled=512, num_classes=self.vocab_size)
                losses.append(step_loss)
                outputs.append(tf.matmul(rnn_t, tf.transpose(self.w_proj)) + self.b_proj)
            self.step_losses = losses
            self.output = outputs
            self.loss = tf.reduce_mean(self.step_losses)
            self.softmax = tf.nn.softmax(self.output)