def __init__(self, args, is_training=True, batch=True): self.args = args if not is_training: args.seq_length = 1 if not batch: args.batch_size = 1 if args.model == 'rnn': cell = rnn_cell.BasicRNNCell(args.rnn_size) elif args.model == 'gru': cell = rnn_cell.GRUCell(args.rnn_size) elif args.model == 'lstm': cell = rnn_cell.BasicLSTMCell(args.rnn_size) else: raise Exception("model type not supported: {}".format(args.model)) self.cell = rnn_cell.MultiRNNCell([cell] * args.num_layers) self.input_data = tf.placeholder(tf.int32, [args.batch_size, args.seq_length]) self.targets = tf.placeholder(tf.int32, [args.batch_size, args.seq_length]) self.initial_state = self.cell.zero_state(args.batch_size, tf.float32) with tf.variable_scope('rnn'): softmax_w = tf.get_variable('softmax_w', [args.rnn_size, args.vocab_size]) softmax_b = tf.get_variable('softmax_b', [args.vocab_size]) with tf.device('/cpu:0'): embedding = tf.get_variable('embedding', [args.vocab_size, args.rnn_size]) inputs = tf.split(1, args.seq_length, tf.nn.embedding_lookup(embedding, self.input_data)) inputs = [tf.squeeze(i, [1]) for i in inputs] def loop(prev, _): prev = tf.nn.xw_plus_b(prev, softmax_w, softmax_b) prev_symbol = tf.stop_gradient(tf.argmax(prev, 1)) return tf.nn.embedding_lookup(embedding, prev_symbol) outputs, last_state = seq2seq.rnn_decoder(inputs, self.initial_state, self.cell, loop_function=None if is_training else loop, scope='rnn') output = tf.reshape(tf.concat(1, outputs), [-1, args.rnn_size]) self.logits = tf.nn.xw_plus_b(output, softmax_w, softmax_b) self.probs = tf.nn.softmax(self.logits) loss = seq2seq.sequence_loss_by_example([self.logits], [tf.reshape(self.targets, [-1])], [tf.ones([args.batch_size * args.seq_length])], args.vocab_size) self.cost = tf.reduce_sum(loss) / args.batch_size / args.seq_length self.final_state = last_state self.lr = tf.Variable(0.0, trainable = False) tvars = tf.trainable_variables() grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars, aggregation_method=2), args.grad_clip) optimizer = tf.train.AdamOptimizer(self.lr) self.train_op = optimizer.apply_gradients(zip(grads, tvars))
def __init__(self, args, is_training=True): self.args = args if not is_training: args.batch_size = 1 args.seq_length = 1 if args.model == 'rnn': self.cell = rnn_cell.BasicRNNCell(args.rnn_size) elif args.model == 'gru': self.cell = rnn_cell.GRUCell(args.rnn_size) elif args.model == 'lstm': self.cell = rnn_cell.BasicLSTMCell(args.rnn_size) else: raise Exception('model type not supported: {}'.format(args.model)) self.cell = rnn_cell.MultiRNNCell([self.cell] * args.num_layers) self.input_data = tf.placeholder(tf.int32, [args.batch_size, args.seq_length]) self.targets = tf.placeholder(tf.int32, [args.batch_size, args.seq_length]) # Target replication self.initial_state = self.cell.zero_state(args.batch_size, tf.float32) with tf.variable_scope('rnn'): softmax_w = tf.get_variable('softmax_w', [args.rnn_size, 2]) softmax_b = tf.get_variable('softmax_b', [2]) with tf.device('/cpu:0'): embedding = tf.get_variable('embedding', [args.vocab_size, args.rnn_size]) inputs = tf.split(1, args.seq_length, tf.nn.embedding_lookup(embedding, self.input_data)) inputs = [tf.squeeze(i, [1]) for i in inputs] outputs, last_state = seq2seq.rnn_decoder(inputs, self.initial_state, self.cell, loop_function=None) output_tf = tf.reshape(tf.concat(1, outputs), [-1, args.rnn_size]) self.logits = tf.nn.xw_plus_b(output_tf, softmax_w, softmax_b) self.probs = tf.nn.softmax(self.logits) loss = seq2seq.sequence_loss_by_example( [self.logits], [tf.reshape(self.targets, [-1])], [tf.ones([args.batch_size * args.seq_length])]) self.cost = tf.reduce_sum(loss) / args.batch_size / args.seq_length self.final_state = last_state self.lr = tf.Variable(0.0, trainable = False) tvars = tf.trainable_variables() grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars, aggregation_method=2), args.grad_clip) optimizer = tf.train.AdamOptimizer(self.lr) self.train_op = optimizer.apply_gradients(zip(grads, tvars))
def __init__(self, config, is_training=False): self.config = config self.batch_size = batch_size = config.batch_size self.num_steps = num_steps = config.num_steps self.hidden_size = hidden_size = config.hidden_size self.num_layers = 1 vocab_size = config.vocab_size self.max_grad_norm = config.max_grad_norm self.use_lstm = config.use_lstm # Placeholders for inputs. self.input_data = tf.placeholder(tf.int32, [batch_size, num_steps]) self.targets = tf.placeholder(tf.int32, [batch_size, num_steps]) self.initial_state = array_ops.zeros( array_ops.pack([self.batch_size, self.num_steps]), dtype=tf.float32).set_shape([None, self.num_steps]) embedding = tf.get_variable( 'embedding', [self.config.vocab_size, self.config.hidden_size]) # Set up ACT cell and inner rnn-type cell for use inside the ACT cell. with tf.variable_scope("rnn"): if self.use_lstm: inner_cell = rnn_cell.BasicLSTMCell(self.config.hidden_size) else: inner_cell = rnn_cell.GRUCell(self.config.hidden_size) with tf.variable_scope("ACT"): act = ACTCell(self.config.hidden_size, inner_cell, config.epsilon, max_computation=config.max_computation, batch_size=self.batch_size) inputs = tf.nn.embedding_lookup(embedding, self.input_data) inputs = [ tf.squeeze(single_input, [1]) for single_input in tf.split(1, self.config.num_steps, inputs) ] self.outputs, final_state = rnn(act, inputs, dtype=tf.float32) # Softmax to get probability distribution over vocab. output = tf.reshape(tf.concat(1, self.outputs), [-1, hidden_size]) softmax_w = tf.get_variable("softmax_w", [hidden_size, vocab_size]) softmax_b = tf.get_variable("softmax_b", [vocab_size]) self.logits = tf.matmul( output, softmax_w) + softmax_b # dim (numsteps*batchsize, vocabsize) loss = seq2seq.sequence_loss_by_example( [self.logits], [tf.reshape(self.targets, [-1])], [tf.ones([batch_size * num_steps])], vocab_size) # Add up loss and retrieve batch-normalised ponder cost: sum N + sum Remainder. ponder_cost = act.calculate_ponder_cost( time_penalty=self.config.ponder_time_penalty) self.cost = (tf.reduce_sum(loss) / batch_size) + ponder_cost self.final_state = self.outputs[-1] if is_training: self.lr = tf.Variable(0.0, trainable=False) tvars = tf.trainable_variables() grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars), self.max_grad_norm) optimizer = tf.train.AdamOptimizer(self.config.learning_rate) self.train_op = optimizer.apply_gradients(zip(grads, tvars))
def __init__(self, args, is_training=True): if not is_training: seq_length = 1 else: seq_length = args.seq_length if args.model == 'rnn': cell_gen = rnn_cell.BasicRNNCell(args.rnn_size) cell_dis = rnn_cell.BasicRNNCell(args.rnn_size) elif args.model == 'gru': cell_gen = rnn_cell.GRUCell(args.rnn_size) cell_dis = rnn_cell.GRUCell(args.rnn_size) elif args.model == 'lstm': cell_gen = rnn_cell.BasicLSTMCell(args.rnn_size) cell_dis = rnn_cell.BasicLSTMCell(args.rnn_size) else: raise Exception('model type not supported: {}'.format(args.model)) # Pass the generated sequences and targets (1) with tf.name_scope('input'): with tf.name_scope('data'): self.input_data = tf.placeholder(tf.int32, [args.batch_size, seq_length]) with tf.name_scope('targets'): self.targets = tf.placeholder(tf.int32, [args.batch_size, seq_length]) ############ # Generator ############ with tf.variable_scope('generator'): self.cell_gen = rnn_cell.MultiRNNCell([cell_gen] * args.num_layers) self.initial_state_gen = self.cell_gen.zero_state(args.batch_size, tf.float32) with tf.variable_scope('rnn'): softmax_w = tf.get_variable('softmax_w', [args.rnn_size, args.vocab_size]) softmax_b = tf.get_variable('softmax_b', [args.vocab_size]) with tf.device('/cpu:0'): embedding = tf.get_variable('embedding', [args.vocab_size, args.rnn_size]) inputs_gen = tf.split(1, seq_length, tf.nn.embedding_lookup( embedding, self.input_data)) inputs_gen = [tf.squeeze(i, [1]) for i in inputs_gen] outputs_gen, last_state_gen = seq2seq.rnn_decoder(inputs_gen, self.initial_state_gen, self.cell_gen, loop_function=None) self.logits_sequence = [] for output_gen in outputs_gen: logits_gen = tf.nn.xw_plus_b(output_gen, softmax_w, softmax_b) self.logits_sequence.append(logits_gen) self.final_state_gen = last_state_gen ################ # Discriminator ################ with tf.variable_scope('discriminator'): self.cell_dis = rnn_cell.MultiRNNCell([cell_dis] * args.num_layers) self.initial_state_dis = self.cell_dis.zero_state(args.batch_size, tf.float32) with tf.variable_scope('rnn'): softmax_w = tf.get_variable('softmax_w', [args.rnn_size, 2]) softmax_b = tf.get_variable('softmax_b', [2]) inputs_dis = [] embedding = tf.get_variable('embedding', [args.vocab_size, args.rnn_size]) for logit in self.logits_sequence: inputs_dis.append(tf.matmul(logit, embedding)) # inputs_dis.append(tf.matmul(tf.nn.softmax(logit), embedding)) outputs_dis, last_state_dis = seq2seq.rnn_decoder(inputs_dis, self.initial_state_dis, self.cell_dis, loop_function=None) probs, logits = [], [] for output_dis in outputs_dis: logit = tf.nn.xw_plus_b(output_dis, softmax_w, softmax_b) prob = tf.nn.softmax(logit) logits.append(logit) probs.append(prob) with tf.name_scope('summary'): probs = tf.pack(probs) probs_real = tf.slice(probs, [0,0,1], [args.seq_length, args.batch_size, 1]) variable_summaries(probs_real, 'probability of real') self.final_state_dis = last_state_dis ######### # Train ######### with tf.name_scope('train'): gen_loss = seq2seq.sequence_loss_by_example( logits, tf.unpack(tf.transpose(self.targets)), tf.unpack(tf.transpose(tf.ones_like(self.targets, dtype=tf.float32)))) self.gen_cost = tf.reduce_sum(gen_loss) / args.batch_size tf.scalar_summary('training loss', self.gen_cost) self.lr_gen = tf.Variable(0.0, trainable = False) self.tvars = tf.trainable_variables() gen_vars = [v for v in self.tvars if not v.name.startswith("discriminator/")] if is_training: gen_grads = tf.gradients(self.gen_cost, gen_vars) self.all_grads = tf.gradients(self.gen_cost, self.tvars) gen_grads_clipped, _ = tf.clip_by_global_norm(gen_grads, args.grad_clip) gen_optimizer = tf.train.AdamOptimizer(self.lr_gen) self.gen_train_op = gen_optimizer.apply_gradients( zip(gen_grads_clipped, gen_vars)) with tf.name_scope('summary'): with tf.name_scope('weight_summary'): for v in self.tvars: variable_summaries(v, v.op.name) if is_training: with tf.name_scope('grad_summary'): for var, grad in zip(self.tvars, self.all_grads): variable_summaries(grad, 'grad/' + var.op.name) self.merged = tf.merge_all_summaries()