import tensorflow as tf
import numpy as np
from config import load_args
from lip_model.losses import cer
from lip_model.modules import embedding, sinusoid_encoding, multihead_attention, \
feedforward, label_smoothing
from lip_model.visual_frontend import VisualFrontend
from util.tf_util import shape_list
config = load_args()
class TransformerTrainGraph():
def __init__(self,
x,
y,
is_training=True,
reuse=None,
embed_input=False,
go_token_index=2,
chars=None):
self.is_training = is_training
self.x = x
if config.featurizer:
vid_inp = x[0] if type(x) is tuple or type(x) is list else x
istarget = tf.not_equal(vid_inp, 0)
self.padding_mask = tf.to_float(tf.reduce_any(istarget, axis=[2,3,4]))
with tf.variable_scope('visual_frontend', reuse=reuse):
self.visual_frontend = VisualFrontend(vid_inp)
vid_inp = self.visual_frontend.output
vid_inp = vid_inp * tf.expand_dims(self.padding_mask,-1)
# pad = 30
# x = tf.keras.layers.ZeroPadding1D(padding=(pad, pad))(x)
if type(x) is tuple or type(x) is list:
x = [vid_inp] + list(x[1:])
else:
x = vid_inp
if is_training:
self.prev = y
self.y = y
else:
# This is the partial prediction used for the autoregression -
# augmented by one more element on every step when autoregression is on
self.prev = y[0]
self.y = y[1] # This is the whole ground truth transcription
self.alignment_history = {} # to be filled in by decoder
self.go_token_idx = go_token_index
# define decoder inputs
self.decoder_inputs = tf.concat(
(tf.ones_like(self.prev[:, :1]) * go_token_index, self.prev[:, :-1]), -1) # 2:<S>
# Encoder
self.enc = x
with tf.variable_scope("encoder", reuse=reuse) as scope:
self.enc = self.encoder_body(self.enc, is_training)
# import ipdb; ipdb.set_trace()
# Decoder
self.dec = self.decoder_inputs
top_scope = tf.get_variable_scope() # this is a hack to be able to use same model
self.chars = chars # needed for decoding with external LM
# --------------- index to char dict for summaries --------------------------------
if chars is not None:
keys = tf.constant( np.arange(len(chars)) , dtype=tf.int64)
values = tf.constant(chars , dtype=tf.string)
self.char_table = tf.contrib.lookup.HashTable(
tf.contrib.lookup.KeyValueTensorInitializer(keys, values), '')
with tf.variable_scope("decoder", reuse=reuse) as scope:
self.dec = self.decoder_body(self.enc, self.dec, is_training, top_scope=top_scope)
if type(self.dec) == tuple:
self.preds, self.scores, self.dec = self.dec # Inference graph output
self.add_loss_and_metrics(reuse, is_training)
if config.tb_eval:
self.add_tb_summaries()
self.tb_sum = tf.summary.merge_all()
def project_output(self):
return True
def decoder_body(self, enc, dec, is_training, top_scope=None):
# Initialize the masks for the pads from here,
# because after positional embeddings are added, nothing will be 0
if config.mask_pads: # Guard this for backwards compatibility
key_masks_enc = tf.sign(tf.abs(tf.reduce_sum(enc, axis=-1))) # (N, T_k)
key_masks_dec = tf.cast( tf.sign(tf.abs(dec)), 'float32' ) # (N, T_k)
query_masks_dec = tf.cast( tf.sign(tf.abs(dec)), 'float32' ) # (N, T_k)
else:
key_masks_enc = key_masks_dec = query_masks_dec = None
## Embedding
dec = self.decoder_embeddings(dec, is_training)
for i in range(config.num_blocks):
with tf.variable_scope("num_blocks_{}".format(i)):
## self-attention
dec, alignmets = multihead_attention(queries=dec,
query_masks=query_masks_dec,
keys=dec,
key_masks=key_masks_dec,
num_units=config.hidden_units,
num_heads=config.num_heads,
dropout_rate=config.dropout_rate,
is_training=is_training,
causality=True,
scope="self_attention")
# self.alignment_history["dec_self_att_{}".format(i)] = alignmets # save for tb
## vanilla attention
dec, alignmets = multihead_attention(queries=dec,
query_masks=query_masks_dec,
keys=enc,
key_masks=key_masks_enc,
num_units=config.hidden_units,
num_heads=config.num_heads,
dropout_rate=config.dropout_rate,
is_training=is_training,
causality=False,
scope="vanilla_attention")
self.alignment_history["enc_dec_attention_{}".format(i)] = alignmets # save for tb
## Feed Forward
dec = feedforward(dec, num_units=[4 * config.hidden_units, config.hidden_units])
return dec
def decoder_embeddings(self, decoder_inputs, is_training):
dec = embedding(decoder_inputs,
vocab_size=config.n_labels,
num_units=config.hidden_units,
scale=True,
scope="dec_embed")
# if self.is_training:
# dec = dec[:,:self.out_last_non_pad_idx]
## Positional Encoding
pos = self.positional_encoding(decoder_inputs, scope='dec_pe')
# if self.is_training:
# pos = pos[:,:self.out_last_non_pad_idx]
dec += pos
## Dropout
dec = tf.layers.dropout(dec,
rate=config.dropout_rate,
training=tf.convert_to_tensor(is_training))
return dec
def positional_encoding(self, inp, scope):
if config.sinusoid:
return sinusoid_encoding(inp,
num_units=config.hidden_units,
zero_pad=False,
scale=False,
scope=scope,
T = config.maxlen
)
else:
return embedding(
tf.tile(tf.expand_dims(tf.range(tf.shape(inp)[1]), 0),
[tf.shape(inp)[0], 1]),
vocab_size=config.maxlen,
num_units=config.hidden_units,
zero_pad=False,
scale=False,
scope="dec_pe")
def encoder_body(self, enc, is_training):
num_blocks = config.num_blocks
if config.mask_pads: # Guard this for backwards compatibility
# Initialize the masks for the pads from here,
# because after positional embeddings are added, nothing will be 0
key_masks = tf.sign(tf.abs(tf.reduce_sum(enc, axis=-1))) # (N, T_k)
query_masks = tf.sign(tf.abs(tf.reduce_sum(enc, axis=-1))) # (N, T_k)
else:
key_masks = query_masks = None
enc = self.encoder_embeddings(enc, is_training)
for i in range(num_blocks):
with tf.variable_scope("num_blocks_{}".format(i)):
### Multihead Attention
enc, alignmets = multihead_attention(queries=enc,
query_masks=query_masks,
keys=enc,
key_masks=key_masks,
num_units=config.hidden_units,
num_heads=config.num_heads,
dropout_rate=config.dropout_rate,
is_training=is_training,
causality=False)
# key_masks = query_masks = None #
### Feed Forward
enc = feedforward(enc, num_units=[4 * config.hidden_units, config.hidden_units])
# self.alignment_history["enc_self_att_{}".format(i)] = alignmets # save for tb
return enc
def encoder_embeddings(self, x, is_training, embed_input=0):
# Embedding
if embed_input:
enc = embedding(x,
vocab_size=config.n_input_vocab,
num_units=config.hidden_units,
scale=True,
scope="enc_embed")
else:
enc = x
## Positional Encoding
feat_dim = shape_list(enc)[-1]
# if input features are not same size as transformer units, make a linear projection
if not feat_dim == config.hidden_units:
enc = tf.layers.dense(enc, config.hidden_units)
enc += self.positional_encoding(enc, scope='enc_pe')
if embed_input:
## Dropout
enc = tf.layers.dropout(enc,
rate=config.dropout_rate,
training=tf.convert_to_tensor(is_training))
return enc
def add_loss_and_metrics(self, reuse, is_training):
# Final linear projection
if self.project_output():
# with tf.variable_scope("ctc_conv1d_net/ctc_probs", reuse=reuse) as scope:
self.logits = tf.layers.dense(self.dec, config.n_labels, reuse=reuse)
else:
assert(self.dec.get_shape()[-1].value == config.n_labels)
self.logits = self.dec
if config.test_aug_times:
self.logits_aug = self.logits
self.logits = tf.reduce_mean(self.logits, 0, keep_dims=True)
self.istarget = tf.to_float(tf.not_equal(self.y, 0))
# Loss
self.y_one_hot = tf.one_hot(self.y, depth=config.n_labels)
self.y_smoothed = label_smoothing(self.y_one_hot)
self.loss = tf.nn.softmax_cross_entropy_with_logits(logits=self.logits,
labels=self.y_smoothed)
# we want to know when to stop so learn padding as well
self.mean_loss = tf.reduce_sum(self.loss) / (tf.reduce_sum(self.istarget))
self.logprobs = tf.log(tf.nn.softmax(self.logits))
if not 'infer' in config.graph_type:
self.preds = tf.to_int32(tf.argmax(self.logits, axis=-1))
self.cer, self.cer_per_sample = cer(self.y_one_hot, self.logits, return_all=True)
else:
self.preds = tf.to_int32(self.preds)
one_hot_from_preds = tf.one_hot(self.preds, depth=config.n_labels)
self.cer, self.cer_per_sample = cer(self.y_one_hot,
one_hot_from_preds,
return_all=True)
def add_tb_summaries(self):
from util.tb_util import add_gif_summary, colorize_image
fps = 10
timeline = False
# ---------------- Add video summaries -------------------------------
bs = int(self.visual_frontend.output.shape[0])
b_id = 0
non_pad_inds = tf.cast(tf.where(self.padding_mask[b_id] > 0)[:, 0], tf.int64)
fr_in, to_in = non_pad_inds[0], non_pad_inds[-1] + 1 # For masking out input paddings
add_gif_summary('1-video_input',
self.visual_frontend.input[b_id][fr_in:to_in], fps=fps, timeline=timeline)
if not config.test_aug_times:
add_gif_summary('2-input_to_resnet',
self.visual_frontend.aug_out[b_id][fr_in:to_in], fps=fps, timeline=timeline)
else:
# Viz the different test augmentations
add_gif_summary('2-input_to_resnet',
tf.concat([ self.visual_frontend.aug_out[b_id][fr_in:to_in]
for b_id in xrange(bs) ], axis=2), fps=fps, timeline=timeline)
# ---------------- Add text summaries -------------------------------
pred_strings_tf = self.char_table.lookup(tf.cast(self.preds, tf.int64))
joined_pred = tf.string_join(
tf.split(pred_strings_tf, pred_strings_tf.shape[1], 1))[ :, 0]
gt_strings_tf = self.char_table.lookup(tf.cast(self.y, tf.int64))
joined_gt = tf.string_join(
tf.split(gt_strings_tf, pred_strings_tf.shape[1], 1))[:, 0]
joined_all = tf.string_join([joined_gt, joined_pred], ' --> ')
tf.summary.text('Predictions', joined_all)
# ---------------- Add image summaries -------------------------------
all_atts = []
for layer_name, alignment_history in self.alignment_history.items():
for att_head_idx, attention_images in enumerate(alignment_history):
all_atts.append(attention_images)
avg_att = tf.exp(tf.reduce_mean(tf.log(all_atts), axis=0))
# Permute and reshape (batch, t_dec, t_enc) --> (batch, t_enc, t_dec, 1)
attention_img = tf.expand_dims(tf.transpose(avg_att, [0, 2, 1]), -1)
attention_img *= 255 # Scale to range [0, 255]
b_id = 0 # visualize only the first sample of the batch
to_out = tf.where( self.preds[b_id]> 0 )[-1][0] + 1 # To mask output paddings |~
color_img = tf.map_fn( colorize_image, (attention_img[:, fr_in:to_in, :to_out]) )
tf.summary.image("3-enc_dec_attention", color_img)
# ---------------- Add image with subs summaries -------------------------------
# import ipdb; ipdb.set_trace()
add_gif_summary('4-subs',
self.visual_frontend.input[b_id][fr_in:to_in], fps=fps, timeline=timeline,
attention=attention_img[b_id][fr_in:to_in, :to_out,0], pred=joined_pred[b_id])
@classmethod
def get_input_shapes_and_types(cls, batch=0):
input_types = []
input_shape = []
if config.featurizer:
input_shape += [ (config.time_dim, config.img_width, config.img_height, config.img_channels) ]
input_types +=['float32']
else:
input_shape += [ (config.time_dim, config.feat_dim) ]
input_types +=['float32']
if batch:
input_shape = [ (config.batch_size,) + shape for shape in input_shape ]
return input_shape, input_types
@classmethod
def get_target_shapes_and_types(cls, batch=0):
target_shape = [ ( config.time_dim, ) ]
if batch:
target_shape = [ (config.batch_size,) + shape for shape in target_shape ]
target_types = ['int64']
return target_shape, target_types
@classmethod
def get_model_input_target_shapes_and_types(cls, batch_dims=1):
return cls.get_input_shapes_and_types(batch=batch_dims),\
cls.get_target_shapes_and_types(batch=batch_dims)
