Created by: OleNet

可以直接在本地跑并且复现错误的代码：

import math
import numpy as np
import paddle
import paddle.fluid as fluid
import sys
import unittest

def kl_divergence(q_logits, p_logits, regression=False):
    if regression:
        q = fluid.layers.sigmoid(q_logits)
        kl = (fluid.layers.sigmoid_cross_entropy_with_logits(x=p_logits, label=q) \
            - fluid.layers.sigmoid_cross_entropy_with_logits(x=q_logits, label=q))
    else:
        q = fluid.layers.softmax(q_logits)
        p = fluid.layers.softmax(p_logits)
        kl = fluid.layers.reduce_sum(q * (fluid.layers.log(q) - fluid.layers.log(p)), -1)

    loss = fluid.layers.mean(x=kl)
    return loss


def scale_l2(x, norm_length):   
    alpha = fluid.layers.reduce_max(fluid.layers.abs(x), dim=1, keep_dim=True) + 1e-12
    l2_norm = alpha * fluid.layers.sqrt(
            fluid.layers.reduce_sum(fluid.layers.pow(x / alpha, factor=2), dim=1, keep_dim=True) + 1e-6) + 1e-12
    x_unit = x / l2_norm
    return norm_length * x_unit


def vatm(logits, 
         embedded, 
         logits_from_embedding_fn, 
         k_iteration=8, 
         small_constant_for_finite_diff=1e-5, 
         return_emb=False):

    program = fluid.default_main_program()

    d = fluid.layers.gaussian_random(embedded.shape)
    
    for _ in range(k_iteration):
        d = scale_l2(d, small_constant_for_finite_diff)
        d_logtis = logits_from_embedding_fn(embedded + d)
        kl = kl_divergence(logits, d_logtis)
        param_grads = fluid.backward.append_backward(kl, parameter_list=[d.name])
        
        gradient = list(filter(lambda p: p[0].name == d.name, param_grads))[0][1]
        gradient.stop_gradient = True
        d = gradient

    d = scale_l2(d, small_constant_for_finite_diff)
    d_logtis = logits_from_embedding_fn(embedded + d)
    kl = kl_divergence(logits, d_logtis)

    if return_emb == False:
        return kl
    else:
        return kl, embedded + d
    

def exe_run(mp, sp, feed, fetch_list):
    place = fluid.CPUPlace()
    exe = fluid.Executor(place)
    exe.run(sp)
    return exe.run(mp, 
            feed=feed,
            fetch_list=fetch_list)


class DummyTextModel(object):
    def __init__(self, nb_classes=10):
        self.ids = fluid.layers.data(name='ids', shape=[1], dtype='int64')
        self.emb = fluid.layers.embedding(
            input=self.ids,
            size=[10, 3],
            dtype="float32")
    
    def forward(self):
        mp, sp = fluid.Program(), fluid.Program()
        return self.forward_from_embedding(self.emb)

    def forward_from_embedding(self, emb):
        net = paddle.fluid.layers.reduce_mean(emb, dim=1)
        net = paddle.fluid.layers.flatten(net)
        net = paddle.fluid.layers.fc(net, 60)
        logits = paddle.fluid.layers.fc(net, 10)
        return paddle.fluid.layers.softmax(logits)
    
    
mp, sp = fluid.Program(), fluid.Program()

with fluid.program_guard(mp, sp):
    text_model = DummyTextModel()
    logits = text_model.forward()

    logits_from_embedding_fn = text_model.forward_from_embedding

    _, emb = vatm(logits,
                          text_model.emb,
                          logits_from_embedding_fn,
                          k_iteration=8, 
                          small_constant_for_finite_diff=1e-5, 
                          return_emb=True)


exe_run(mp, sp, feed={'ids':[[[1], [3]]]}, fetch_list=[])