infer_network.py

# -*- coding=utf-8 -*-
"""
#   Copyright (c) 2020 PaddlePaddle Authors. 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.
"""

import math
import argparse
import numpy as np
import paddle.fluid as fluid
from utils import tdm_sampler_prepare, tdm_child_prepare, trace_var
from train_network import DnnLayerClassifierNet, InputTransNet


class TdmInferNet(object):
    def __init__(self, args):
        self.input_embed_size = args.query_emb_size
        self.node_embed_size = args.node_emb_size
        self.label_nums = 2  # label为正负两类
        self.node_nums = args.node_nums
        self.max_layers = args.layer_size
        self.batch_size = args.batch_size
        self.topK = args.topK  # 最终召回多少个item
        self.child_nums = args.child_nums  # 若树为二叉树，则child_nums=2

        self.layer_list = self.get_layer_list(args)
        self.first_layer_idx = 0
        self.first_layer_node = self.create_first_layer(args)
        self.layer_classifier = DnnLayerClassifierNet(args)
        self.input_trans_net = InputTransNet(args)

    def input_data(self):
        input_emb = fluid.layers.data(
            name="input_emb",
            shape=[self.input_embed_size],
            dtype="float32",
        )

        # first_layer 与 first_layer_mask 对应着infer起始层的节点
        first_layer = fluid.layers.data(
            name="first_layer_node",
            shape=[1],
            dtype="int64",
            lod_level=1,
        )

        first_layer_mask = fluid.layers.data(
            name="first_layer_node_mask",
            shape=[1],
            dtype="int64",
            lod_level=1,
        )

        inputs = [input_emb] + [first_layer] + [first_layer_mask]
        return inputs

    def get_layer_list(self, args):
        """get layer list from layer_list.txt"""
        layer_list = []
        with open(args.tree_layer_init_path, 'r') as fin:
            for line in fin.readlines():
                l = []
                layer = (line.split('\n'))[0].split(',')
                for node in layer:
                    if node:
                        l.append(node)
                layer_list.append(l)
        return layer_list

    def create_first_layer(self, args):
        """decide which layer to start infer"""
        first_layer_id = 0
        for idx, layer_node in enumerate(args.layer_node_num_list):
            if layer_node >= self.topK:
                first_layer_id = idx
                break
        first_layer_node = self.layer_list[first_layer_id]
        self.first_layer_idx = first_layer_id
        return first_layer_node

    def infer_net(self, inputs):
        """
        infer的主要流程
        infer的基本逻辑是：从上层开始（具体层idx由树结构及TopK值决定）
        1、依次通过每一层分类器，得到当前层输入的指定节点的prob
        2、根据prob值大小，取topK的节点，取这些节点的孩子节点作为下一层的输入
        3、循环1、2步骤，遍历完所有层，得到每一层筛选结果的集合
        4、将筛选结果集合中的叶子节点，拿出来再做一次topK，得到最终的召回输出
        """
        input_emb = inputs[0]
        current_layer_node = inputs[1]
        current_layer_child_mask = inputs[2]

        node_score = []
        node_list = []

        input_trans_emb = self.input_trans_net.input_fc_infer(input_emb)

        for layer_idx in range(self.first_layer_idx, self.max_layers):
            # 确定当前层的需要计算的节点数
            if layer_idx == self.first_layer_idx:
                current_layer_node_num = len(self.first_layer_node)
            else:
                current_layer_node_num = current_layer_node.shape[1] * \
                    current_layer_node.shape[2]

            current_layer_node = fluid.layers.reshape(
                current_layer_node, [-1, current_layer_node_num])
            current_layer_child_mask = fluid.layers.reshape(
                current_layer_child_mask, [-1, current_layer_node_num])

            node_emb = fluid.embedding(
                input=current_layer_node,
                size=[self.node_nums, self.node_embed_size],
                param_attr=fluid.ParamAttr(name="TDM_Tree_Emb"))

            input_fc_out = self.input_trans_net.layer_fc_infer(
                input_trans_emb, layer_idx)

            # 过每一层的分类器
            layer_classifier_res = self.layer_classifier.classifier_layer_infer(input_fc_out,
                                                                                node_emb,
                                                                                layer_idx)

            # 过最终的判别分类器
            tdm_fc = fluid.layers.fc(input=layer_classifier_res,
                                     size=self.label_nums,
                                     act=None,
                                     num_flatten_dims=2,
                                     param_attr=fluid.ParamAttr(
                                         name="tdm.cls_fc.weight"),
                                     bias_attr=fluid.ParamAttr(name="tdm.cls_fc.bias"))

            prob = fluid.layers.softmax(tdm_fc)
            positive_prob = fluid.layers.slice(
                prob, axes=[2], starts=[1], ends=[2])
            prob_re = fluid.layers.reshape(
                positive_prob, [-1, current_layer_node_num])

            # 过滤掉padding产生的无效节点（node_id=0）
            node_zero_mask = fluid.layers.cast(current_layer_node, 'bool')
            node_zero_mask = fluid.layers.cast(node_zero_mask, 'float')
            prob_re = prob_re * node_zero_mask

            # 在当前层的分类结果中取topK，并将对应的score及node_id保存下来
            k = self.topK
            if current_layer_node_num < self.topK:
                k = current_layer_node_num
            _, topk_i = fluid.layers.topk(prob_re, k)

            # index_sample op根据下标索引tensor对应位置的值
            # 若paddle版本>2.0，调用方式为paddle.index_sample
            top_node = fluid.contrib.layers.index_sample(
                current_layer_node, topk_i)
            prob_re_mask = prob_re * current_layer_child_mask  # 过滤掉非叶子节点
            topk_value = fluid.contrib.layers.index_sample(
                prob_re_mask, topk_i)
            node_score.append(topk_value)
            node_list.append(top_node)

            # 取当前层topK结果的孩子节点，作为下一层的输入
            if layer_idx < self.max_layers - 1:
                # tdm_child op 根据输入返回其 child 及 child_mask
                # 若child是叶子节点，则child_mask=1，否则为0
                current_layer_node, current_layer_child_mask = \
                    fluid.contrib.layers.tdm_child(x=top_node,
                                                   node_nums=self.node_nums,
                                                   child_nums=self.child_nums,
                                                   param_attr=fluid.ParamAttr(
                                                       name="TDM_Tree_Info"),
                                                   dtype='int64')

        total_node_score = fluid.layers.concat(node_score, axis=1)
        total_node = fluid.layers.concat(node_list, axis=1)

        # 考虑到树可能是不平衡的，计算所有层的叶子节点的topK
        res_score, res_i = fluid.layers.topk(total_node_score, self.topK)
        res_layer_node = fluid.contrib.layers.index_sample(total_node, res_i)
        res_node = fluid.layers.reshape(res_layer_node, [-1, self.topK, 1])

        # 利用Tree_info信息，将node_id转换为item_id
        tree_info = fluid.default_main_program().global_block().var("TDM_Tree_Info")
        res_node_emb = fluid.layers.gather_nd(tree_info, res_node)

        res_item = fluid.layers.slice(
            res_node_emb, axes=[2], starts=[0], ends=[1])
        res_item_re = fluid.layers.reshape(res_item, [-1, self.topK])
        return res_item_re