tensorflow 中使用 tf.py_func 灵活地自定义张量计算

背景是这样的，在看完论文《A Deep Reinforced Model for Abstractive Summarization》之后想用 tensorflow 实现一下。论文中的一个关键点是根据摘要生成的一个评测指标 ROUGE 作为强化学习的 reward，参与到损失函数的计算中，需要优化的目标函数如下（详细一些的介绍可参考我的另一篇文章https://www.jianshu.com/p/38a2d3e04272）：

强化学习目标函数.png

去 github 找了一下这个方法的开源实现，发现了两种实现方式。

方式一(有问题）

https://github.com/weili-ict/SelfCriticalSequenceTraining-tensorflow
此方式的实现思路是先通过一次 sess.run 根据最大概率生成序列作为 baseline 以及进行采样得到序列 y^s。生成序列后计算 baseline 和 y^s 的 reward。然后将 reward 放到 placeholder 中再跑一次 sess.run 计算 L_rl 并优化。此方法的问题是两次 sess.run 产生的 baseline 序列是不变的，但 y^s 序列是根据概率分布采样生成的两次生成的序列完全不同，不能将第一次计算出的 reward 用于第二次的损失函数计算。
主要逻辑如下：

captions_batch = np.array(captions[i * self.batch_size:(i + 1) * self.batch_size])
image_idxs_batch = np.array(image_idxs[i * self.batch_size:(i + 1) * self.batch_size])
features_batch = np.array(features[image_idxs_batch])

ground_truths = [captions[image_idxs == image_idxs_batch[j]] for j in
                 range(len(image_idxs_batch))]
ref_decoded = [decode_captions(ground_truths[j], self.model.idx_to_word) for j in range(len(ground_truths))]

feed_dict = {self.model.features: features_batch, self.model.captions: captions_batch}
# first run to get 2 different serials
samples, greedy_words = sess.run([sampled_captions, greedy_caption],
                                         feed_dict)
masks, all_decoded = decode_captions_for_blue(samples, self.model.idx_to_word)
_, greedy_decoded = decode_captions_for_blue(greedy_words, self.model.idx_to_word)
# calculate the rewards of 2 serials
r = [evaluate_captions([k], [v])  for k, v in zip(ref_decoded, all_decoded)]
b = [evaluate_captions([k], [v]) for k, v in zip(ref_decoded, greedy_decoded)]

b_for_eval.extend(b)
feed_dict = {grad_mask: masks, rewards: r, base_line: b,
             self.model.features: features_batch, self.model.captions: captions_batch
             } 
# calculate loss and train
_ = sess.run([train_op], feed_dict)

方式二

https://github.com/ne7ermore/deeping-flow/tree/master/deep-reinforced-sum-model
由方式一可见，进行一次序列生成 -> 计算序列的 reward -> 进行第二次序列生成，利用已计算的 reward 计算 loss 并优化，这一模式是行不通的，由于随机采样的的存在，必须保证 reward 计算与 loss 计算优化在同一次 sess.run 中进行。
但是，困难在于 reward 的计算是一个较为复杂的过程， tensorflow 中必然没有提供这样的计算 API，那么如何在 sess.run 中进行这样复杂的自定义的计算呢？
可以利用 tf.py_func 函数。py_func 函数接收一个用户自定义的函数 f 和一个输入张量 input 作为参数，返回一个利用 f 对 input 转换后的输出张量。
主要逻辑如下：

b_words = model.sample()
s_words, props = model.sample(False)

s_props = tf.reshape(_gather_index(tf.reshape(
    props, [-1, args.tgt_vs]), s_words, model.prev), [args.batch_size, args.l_max_len])
# calculate rewards in session
baseline = tf.py_func(rouge_l, [b_words, model.tgt], tf.float32)
reward = tf.py_func(rouge_l, [s_words, model.tgt], tf.float32)
advantage = reward - baseline

mask = pad_mask(model.tgt, EOS, [args.batch_size, args.l_max_len])

loss = -tf.reduce_sum(s_props *
                      mask * advantage[:, None]) / tf.reduce_sum(mask)

可以看到此方式将序列生成、reward 计算、loss 计算放在同一个 sess.run 中进行，解决了方式一中两次 sess.run 生成的采样候选序列不同的问题。

参考：
1、《A Deep Reinforced Model for Abstractive Summarization》
2、《Self-critical Sequence Training for Image Captioning》
3、https://github.com/weili-ict/SelfCriticalSequenceTraining-tensorflow
4、https://github.com/ne7ermore/deeping-flow/tree/master/deep-reinforced-sum-model