research/adversarial_text/layers.py
# Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Layers for VatxtModel."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# Dependency imports
from six.moves import xrange
import tensorflow as tf
K = tf.keras
def cl_logits_subgraph(layer_sizes, input_size, num_classes, keep_prob=1.):
"""Construct multiple ReLU layers with dropout and a linear layer."""
subgraph = K.models.Sequential(name='cl_logits')
for i, layer_size in enumerate(layer_sizes):
if i == 0:
subgraph.add(
K.layers.Dense(layer_size, activation='relu', input_dim=input_size))
else:
subgraph.add(K.layers.Dense(layer_size, activation='relu'))
if keep_prob < 1.:
subgraph.add(K.layers.Dropout(1. - keep_prob))
subgraph.add(K.layers.Dense(1 if num_classes == 2 else num_classes))
return subgraph
class Embedding(K.layers.Layer):
"""Embedding layer with frequency-based normalization and dropout."""
def __init__(self,
vocab_size,
embedding_dim,
normalize=False,
vocab_freqs=None,
keep_prob=1.,
**kwargs):
self.vocab_size = vocab_size
self.embedding_dim = embedding_dim
self.normalized = normalize
self.keep_prob = keep_prob
if normalize:
assert vocab_freqs is not None
self.vocab_freqs = tf.constant(
vocab_freqs, dtype=tf.float32, shape=(vocab_size, 1))
super(Embedding, self).__init__(**kwargs)
def build(self, input_shape):
with tf.device('/cpu:0'):
self.var = self.add_weight(
shape=(self.vocab_size, self.embedding_dim),
initializer=tf.random_uniform_initializer(-1., 1.),
name='embedding',
dtype=tf.float32)
if self.normalized:
self.var = self._normalize(self.var)
super(Embedding, self).build(input_shape)
def call(self, x):
embedded = tf.nn.embedding_lookup(self.var, x)
if self.keep_prob < 1.:
shape = embedded.get_shape().as_list()
# Use same dropout masks at each timestep with specifying noise_shape.
# This slightly improves performance.
# Please see https://arxiv.org/abs/1512.05287 for the theoretical
# explanation.
embedded = tf.nn.dropout(
embedded, self.keep_prob, noise_shape=(shape[0], 1, shape[2]))
return embedded
def _normalize(self, emb):
weights = self.vocab_freqs / tf.reduce_sum(self.vocab_freqs)
mean = tf.reduce_sum(weights * emb, 0, keep_dims=True)
var = tf.reduce_sum(weights * tf.pow(emb - mean, 2.), 0, keep_dims=True)
stddev = tf.sqrt(1e-6 + var)
return (emb - mean) / stddev
class LSTM(object):
"""LSTM layer using dynamic_rnn.
Exposes variables in `trainable_weights` property.
"""
def __init__(self, cell_size, num_layers=1, keep_prob=1., name='LSTM'):
self.cell_size = cell_size
self.num_layers = num_layers
self.keep_prob = keep_prob
self.reuse = None
self.trainable_weights = None
self.name = name
def __call__(self, x, initial_state, seq_length):
with tf.variable_scope(self.name, reuse=self.reuse) as vs:
cell = tf.contrib.rnn.MultiRNNCell([
tf.contrib.rnn.BasicLSTMCell(
self.cell_size,
forget_bias=0.0,
reuse=tf.get_variable_scope().reuse)
for _ in xrange(self.num_layers)
])
# shape(x) = (batch_size, num_timesteps, embedding_dim)
lstm_out, next_state = tf.nn.dynamic_rnn(
cell, x, initial_state=initial_state, sequence_length=seq_length)
# shape(lstm_out) = (batch_size, timesteps, cell_size)
if self.keep_prob < 1.:
lstm_out = tf.nn.dropout(lstm_out, self.keep_prob)
if self.reuse is None:
self.trainable_weights = vs.global_variables()
self.reuse = True
return lstm_out, next_state
class SoftmaxLoss(K.layers.Layer):
"""Softmax xentropy loss with candidate sampling."""
def __init__(self,
vocab_size,
num_candidate_samples=-1,
vocab_freqs=None,
**kwargs):
self.vocab_size = vocab_size
self.num_candidate_samples = num_candidate_samples
self.vocab_freqs = vocab_freqs
super(SoftmaxLoss, self).__init__(**kwargs)
self.multiclass_dense_layer = K.layers.Dense(self.vocab_size)
def build(self, input_shape):
input_shape = input_shape[0].as_list()
with tf.device('/cpu:0'):
self.lin_w = self.add_weight(
shape=(input_shape[-1], self.vocab_size),
name='lm_lin_w',
initializer=K.initializers.glorot_uniform())
self.lin_b = self.add_weight(
shape=(self.vocab_size,),
name='lm_lin_b',
initializer=K.initializers.glorot_uniform())
self.multiclass_dense_layer.build(input_shape)
super(SoftmaxLoss, self).build(input_shape)
def call(self, inputs):
x, labels, weights = inputs
if self.num_candidate_samples > -1:
assert self.vocab_freqs is not None
labels_reshaped = tf.reshape(labels, [-1])
labels_reshaped = tf.expand_dims(labels_reshaped, -1)
sampled = tf.nn.fixed_unigram_candidate_sampler(
true_classes=labels_reshaped,
num_true=1,
num_sampled=self.num_candidate_samples,
unique=True,
range_max=self.vocab_size,
unigrams=self.vocab_freqs)
inputs_reshaped = tf.reshape(x, [-1, int(x.get_shape()[2])])
lm_loss = tf.nn.sampled_softmax_loss(
weights=tf.transpose(self.lin_w),
biases=self.lin_b,
labels=labels_reshaped,
inputs=inputs_reshaped,
num_sampled=self.num_candidate_samples,
num_classes=self.vocab_size,
sampled_values=sampled)
lm_loss = tf.reshape(
lm_loss,
[int(x.get_shape()[0]), int(x.get_shape()[1])])
else:
logits = self.multiclass_dense_layer(x)
lm_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
lm_loss = tf.identity(
tf.reduce_sum(lm_loss * weights) / _num_labels(weights),
name='lm_xentropy_loss')
return lm_loss
def classification_loss(logits, labels, weights):
"""Computes cross entropy loss between logits and labels.
Args:
logits: 2-D [timesteps*batch_size, m] float tensor, where m=1 if
num_classes=2, otherwise m=num_classes.
labels: 1-D [timesteps*batch_size] integer tensor.
weights: 1-D [timesteps*batch_size] float tensor.
Returns:
Loss scalar of type float.
"""
inner_dim = logits.get_shape().as_list()[-1]
with tf.name_scope('classifier_loss'):
# Logistic loss
if inner_dim == 1:
loss = tf.nn.sigmoid_cross_entropy_with_logits(
logits=tf.squeeze(logits, -1), labels=tf.cast(labels, tf.float32))
# Softmax loss
else:
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
num_lab = _num_labels(weights)
tf.summary.scalar('num_labels', num_lab)
return tf.identity(
tf.reduce_sum(weights * loss) / num_lab, name='classification_xentropy')
def accuracy(logits, targets, weights):
"""Computes prediction accuracy.
Args:
logits: 2-D classifier logits [timesteps*batch_size, num_classes]
targets: 1-D [timesteps*batch_size] integer tensor.
weights: 1-D [timesteps*batch_size] float tensor.
Returns:
Accuracy: float scalar.
"""
with tf.name_scope('accuracy'):
eq = tf.cast(tf.equal(predictions(logits), targets), tf.float32)
return tf.identity(
tf.reduce_sum(weights * eq) / _num_labels(weights), name='accuracy')
def predictions(logits):
"""Class prediction from logits."""
inner_dim = logits.get_shape().as_list()[-1]
with tf.name_scope('predictions'):
# For binary classification
if inner_dim == 1:
pred = tf.cast(tf.greater(tf.squeeze(logits, -1), 0.), tf.int64)
# For multi-class classification
else:
pred = tf.argmax(logits, 2)
return pred
def _num_labels(weights):
"""Number of 1's in weights. Returns 1. if 0."""
num_labels = tf.reduce_sum(weights)
num_labels = tf.where(tf.equal(num_labels, 0.), 1., num_labels)
return num_labels
def optimize(loss,
global_step,
max_grad_norm,
lr,
lr_decay,
sync_replicas=False,
replicas_to_aggregate=1,
task_id=0):
"""Builds optimization graph.
* Creates an optimizer, and optionally wraps with SyncReplicasOptimizer
* Computes, clips, and applies gradients
* Maintains moving averages for all trainable variables
* Summarizes variables and gradients
Args:
loss: scalar loss to minimize.
global_step: integer scalar Variable.
max_grad_norm: float scalar. Grads will be clipped to this value.
lr: float scalar, learning rate.
lr_decay: float scalar, learning rate decay rate.
sync_replicas: bool, whether to use SyncReplicasOptimizer.
replicas_to_aggregate: int, number of replicas to aggregate when using
SyncReplicasOptimizer.
task_id: int, id of the current task; used to ensure proper initialization
of SyncReplicasOptimizer.
Returns:
train_op
"""
with tf.name_scope('optimization'):
# Compute gradients.
tvars = tf.trainable_variables()
grads = tf.gradients(
loss,
tvars,
aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
# Clip non-embedding grads
non_embedding_grads_and_vars = [(g, v) for (g, v) in zip(grads, tvars)
if 'embedding' not in v.op.name]
embedding_grads_and_vars = [(g, v) for (g, v) in zip(grads, tvars)
if 'embedding' in v.op.name]
ne_grads, ne_vars = zip(*non_embedding_grads_and_vars)
ne_grads, _ = tf.clip_by_global_norm(ne_grads, max_grad_norm)
non_embedding_grads_and_vars = zip(ne_grads, ne_vars)
grads_and_vars = embedding_grads_and_vars + list(non_embedding_grads_and_vars)
# Summarize
_summarize_vars_and_grads(grads_and_vars)
# Decaying learning rate
lr = tf.train.exponential_decay(
lr, global_step, 1, lr_decay, staircase=True)
tf.summary.scalar('learning_rate', lr)
opt = tf.train.AdamOptimizer(lr)
# Track the moving averages of all trainable variables.
variable_averages = tf.train.ExponentialMovingAverage(0.999, global_step)
# Apply gradients
if sync_replicas:
opt = tf.train.SyncReplicasOptimizer(
opt,
replicas_to_aggregate,
variable_averages=variable_averages,
variables_to_average=tvars,
total_num_replicas=replicas_to_aggregate)
apply_gradient_op = opt.apply_gradients(
grads_and_vars, global_step=global_step)
with tf.control_dependencies([apply_gradient_op]):
train_op = tf.no_op(name='train_op')
# Initialization ops
tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS,
opt.get_chief_queue_runner())
if task_id == 0: # Chief task
local_init_op = opt.chief_init_op
tf.add_to_collection('chief_init_op', opt.get_init_tokens_op())
else:
local_init_op = opt.local_step_init_op
tf.add_to_collection('local_init_op', local_init_op)
tf.add_to_collection('ready_for_local_init_op',
opt.ready_for_local_init_op)
else:
# Non-sync optimizer
apply_gradient_op = opt.apply_gradients(grads_and_vars, global_step)
with tf.control_dependencies([apply_gradient_op]):
train_op = variable_averages.apply(tvars)
return train_op
def _summarize_vars_and_grads(grads_and_vars):
tf.logging.info('Trainable variables:')
tf.logging.info('-' * 60)
for grad, var in grads_and_vars:
tf.logging.info(var)
def tag(name, v=var):
return v.op.name + '_' + name
# Variable summary
mean = tf.reduce_mean(var)
tf.summary.scalar(tag('mean'), mean)
with tf.name_scope(tag('stddev')):
stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
tf.summary.scalar(tag('stddev'), stddev)
tf.summary.scalar(tag('max'), tf.reduce_max(var))
tf.summary.scalar(tag('min'), tf.reduce_min(var))
tf.summary.histogram(tag('histogram'), var)
# Gradient summary
if grad is not None:
if isinstance(grad, tf.IndexedSlices):
grad_values = grad.values
else:
grad_values = grad
tf.summary.histogram(tag('gradient'), grad_values)
tf.summary.scalar(tag('gradient_norm'), tf.global_norm([grad_values]))
else:
tf.logging.info('Var %s has no gradient', var.op.name)