easy_rec/python/tools/feature_selection.py

from __future__ import division
from __future__ import print_function

import json
import os
import sys
from collections import OrderedDict

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.python.framework.meta_graph import read_meta_graph_file

from easy_rec.python.utils import config_util
from easy_rec.python.utils import io_util

if tf.__version__ >= '2.0':
  tf = tf.compat.v1

import matplotlib  # NOQA
matplotlib.use('Agg')  # NOQA
import matplotlib.pyplot as plt  # NOQA

tf.app.flags.DEFINE_string('model_type', 'variational_dropout',
                           'feature selection model type')
tf.app.flags.DEFINE_string('config_path', '',
                           'feature selection model config path')
tf.app.flags.DEFINE_string('checkpoint_path', None,
                           'feature selection model checkpoint path')
tf.app.flags.DEFINE_string('output_dir', '',
                           'feature selection result directory')
tf.app.flags.DEFINE_integer(
    'topk', 100, 'select topk importance features for each feature group')
tf.app.flags.DEFINE_string('fg_path', '', 'fg config path')
tf.app.flags.DEFINE_bool('visualize', False,
                         'visualization feature selection result or not')
FLAGS = tf.app.flags.FLAGS


class VariationalDropoutFS:

  def __init__(self,
               config_path,
               output_dir,
               topk,
               checkpoint_path=None,
               fg_path=None,
               visualize=False):
    self._config_path = config_path
    self._output_dir = output_dir
    self._topk = topk
    if not tf.gfile.Exists(self._output_dir):
      tf.gfile.MakeDirs(self._output_dir)
    self._checkpoint_path = checkpoint_path
    self._fg_path = fg_path
    self._visualize = visualize

  def process(self):
    tf.logging.info('Loading logit_p of VariationalDropout layer ...')
    feature_dim_dropout_p_map, embedding_wise_variational_dropout = self._feature_dim_dropout_ratio(
    )

    feature_importance_map = {}
    for group_name, feature_dim_dropout_p in feature_dim_dropout_p_map.items():
      tf.logging.info('Calculating %s feature importance ...' % group_name)
      feature_importance = self._get_feature_importance(
          feature_dim_dropout_p, embedding_wise_variational_dropout)
      feature_importance_map[group_name] = feature_importance

      tf.logging.info('Dump %s  feature importance to csv ...' % group_name)
      self._dump_to_csv(feature_importance, group_name)

      if self._visualize:
        tf.logging.info('Visualizing %s feature importance ...' % group_name)
        if embedding_wise_variational_dropout:
          self._visualize_embedding_dim_importance(feature_dim_dropout_p)
        self._visualize_feature_importance(feature_importance, group_name)

    tf.logging.info('Processing model config ...')
    self._process_config(feature_importance_map)

  def _feature_dim_dropout_ratio(self):
    """Get dropout ratio of embedding-wise or feature-wise."""
    config = config_util.get_configs_from_pipeline_file(self._config_path)
    assert config.model_config.HasField(
        'variational_dropout'), 'variational_dropout must be in model_config'

    embedding_wise_variational_dropout = config.model_config.variational_dropout.embedding_wise_variational_dropout

    if self._checkpoint_path is None or len(self._checkpoint_path) == 0:
      checkpoint_path = tf.train.latest_checkpoint(config.model_dir)
    else:
      checkpoint_path = self._checkpoint_path

    meta_graph_def = read_meta_graph_file(checkpoint_path + '.meta')
    features_dimension_map = dict()
    for col_def in meta_graph_def.collection_def[
        'variational_dropout'].bytes_list.value:
      name, features_dimension = json.loads(col_def)
      name = 'all' if name == '' else name
      features_dimension_map[name] = OrderedDict(features_dimension)

    tf.logging.info('Reading checkpoint from %s ...' % checkpoint_path)
    reader = tf.train.NewCheckpointReader(checkpoint_path)

    feature_dim_dropout_p_map = {}
    for feature_group in config.model_config.feature_groups:
      group_name = feature_group.group_name

      logit_p_name = 'logit_p' if group_name == 'all' else 'logit_p_%s' % group_name
      try:
        logit_p = reader.get_tensor(logit_p_name)
      except Exception:
        print('get `logit_p` failed, try to get `backbone/logit_p`')
        logit_p = reader.get_tensor('backbone/' + logit_p_name)
      feature_dims_importance = tf.sigmoid(logit_p)
      with tf.Session() as sess:
        feature_dims_importance = feature_dims_importance.eval(session=sess)

      feature_dim_dropout_p = {}
      if embedding_wise_variational_dropout:
        index_end = 0
        for feature_name, feature_dim in features_dimension_map[
            group_name].items():
          index_start = index_end
          index_end = index_start + feature_dim
          feature_dim_dropout_p[feature_name] = feature_dims_importance[
              index_start:index_end]
      else:
        index = 0
        for feature_name in features_dimension_map[group_name].keys():
          feature_dim_dropout_p[feature_name] = feature_dims_importance[index]
          index += 1

      feature_dim_dropout_p_map[group_name] = feature_dim_dropout_p
    return feature_dim_dropout_p_map, embedding_wise_variational_dropout

  def _get_feature_importance(self, feature_dim_dropout_p,
                              embedding_wise_variational_dropout):
    """Calculate feature importance."""
    if embedding_wise_variational_dropout:
      feature_importance = {}
      for item in feature_dim_dropout_p.items():
        dropout_rate_mean = np.mean(item[1])
        feature_importance[item[0]] = dropout_rate_mean
      feature_importance = OrderedDict(
          sorted(feature_importance.items(), key=lambda e: e[1]))
    else:
      feature_importance = OrderedDict(
          sorted(feature_dim_dropout_p.items(), key=lambda e: e[1]))
    return feature_importance

  def _process_config(self, feature_importance_map):
    """Process model config and fg config with feature selection."""
    excluded_features = set()
    for group_name, feature_importance in feature_importance_map.items():
      for i, (feature_name, _) in enumerate(feature_importance.items()):
        if i >= self._topk:
          excluded_features.add(feature_name)

    config = config_util.get_configs_from_pipeline_file(self._config_path)
    # keep sequence features and side-infos
    sequence_features = set()
    for feature_group in config.model_config.feature_groups:
      for sequence_feature in feature_group.sequence_features:
        for seq_att_map in sequence_feature.seq_att_map:
          for key in seq_att_map.key:
            sequence_features.add(key)
          for hist_seq in seq_att_map.hist_seq:
            sequence_features.add(hist_seq)
    # compat with din
    for sequence_feature in config.model_config.seq_att_groups:
      for seq_att_map in sequence_feature.seq_att_map:
        for key in seq_att_map.key:
          sequence_features.add(key)
        for hist_seq in seq_att_map.hist_seq:
          sequence_features.add(hist_seq)
    excluded_features = excluded_features - sequence_features

    feature_configs = []
    for feature_config in config_util.get_compatible_feature_configs(config):
      feature_name = feature_config.feature_name if feature_config.HasField('feature_name') \
          else feature_config.input_names[0]
      if feature_name not in excluded_features:
        feature_configs.append(feature_config)

    if config.feature_configs:
      config.ClearField('feature_configs')
      config.feature_configs.extend(feature_configs)
    else:
      config.feature_config.ClearField('features')
      config.feature_config.features.extend(feature_configs)

    for feature_group in config.model_config.feature_groups:
      feature_names = []
      for feature_name in feature_group.feature_names:
        if feature_name not in excluded_features:
          feature_names.append(feature_name)
      feature_group.ClearField('feature_names')
      feature_group.feature_names.extend(feature_names)
    config_util.save_message(
        config,
        os.path.join(self._output_dir, os.path.basename(self._config_path)))

    if self._fg_path is not None and len(self._fg_path) > 0:
      with tf.gfile.Open(self._fg_path) as f:
        fg_json = json.load(f, object_pairs_hook=OrderedDict)
        features = []
        for feature in fg_json['features']:
          if 'feature_name' in feature:
            if feature['feature_name'] not in excluded_features:
              features.append(feature)
          else:
            features.append(feature)
        fg_json['features'] = features
      with tf.gfile.Open(
          os.path.join(self._output_dir, os.path.basename(self._fg_path)),
          'w') as f:
        json.dump(fg_json, f, indent=4)

  def _dump_to_csv(self, feature_importance, group_name):
    """Dump feature importance data to a csv file."""
    with tf.gfile.Open(
        os.path.join(self._output_dir,
                     'feature_dropout_ratio_%s.csv' % group_name), 'w') as f:
      df = pd.DataFrame(
          columns=['feature_name', 'mean_drop_p'],
          data=[list(kv) for kv in feature_importance.items()])
      df.to_csv(f, encoding='gbk')

  def _visualize_embedding_dim_importance(self, feature_dim_dropout_p):
    """Visualize embedding-wise importance visualization for every feature."""
    output_dir = os.path.join(self._output_dir, 'feature_dims_importance_pics')
    if not tf.gfile.Exists(output_dir):
      tf.gfile.MakeDirs(output_dir)

    plt.rcdefaults()
    for feature_name, feature_dropout_p in feature_dim_dropout_p.items():
      embedding_len = len(feature_dropout_p)
      embedding_dims = []
      for i in range(embedding_len):
        embedding_dims.append('dim_' + str(i + 1))
      y_pos = np.arange(len(embedding_dims))
      performance_list = []
      for i in range(0, embedding_len):
        performance_list.append(feature_dropout_p[i])
      fig, ax = plt.subplots()
      b = ax.barh(
          y_pos,
          performance_list,
          align='center',
          alpha=0.4,
          label='dropout_rate',
          lw=1)
      for rect in b:
        w = rect.get_width()
        ax.text(
            w,
            rect.get_y() + rect.get_height() / 2,
            '%.4f' % w,
            ha='left',
            va='center')
      plt.yticks(y_pos, embedding_dims)
      plt.xlabel(feature_name)
      plt.title('Dropout ratio')
      img_path = os.path.join(output_dir, feature_name + '.png')
      with tf.gfile.GFile(img_path, 'wb') as f:
        plt.savefig(f, format='png')

  def _visualize_feature_importance(self, feature_importance, group_name):
    """Draw feature importance histogram."""
    df = pd.DataFrame(
        columns=['feature_name', 'mean_drop_p'],
        data=[list(kv) for kv in feature_importance.items()])
    df['color'] = ['red' if x < 0.5 else 'green' for x in df['mean_drop_p']]
    df.sort_values('mean_drop_p', inplace=True, ascending=False)
    df.reset_index(inplace=True)
    # Draw plot
    plt.figure(figsize=(90, 200), dpi=100)
    plt.hlines(y=df.index, xmin=0, xmax=df.mean_drop_p)
    for x, y, tex in zip(df.mean_drop_p, df.index, df.mean_drop_p):
      plt.text(
          x,
          y,
          round(tex, 2),
          horizontalalignment='right' if x < 0 else 'left',
          verticalalignment='center',
          fontdict={
              'color': 'red' if x < 0 else 'green',
              'size': 14
          })
    # Decorations
    plt.yticks(df.index, df.feature_name, fontsize=20)
    plt.title('Dropout Ratio', fontdict={'size': 30})
    plt.grid(linestyle='--', alpha=0.5)
    plt.xlim(0, 1)
    with tf.gfile.GFile(
        os.path.join(self._output_dir,
                     'feature_dropout_pic_%s.png' % group_name), 'wb') as f:
      plt.savefig(f, format='png')


if __name__ == '__main__':
  sys.argv = io_util.filter_unknown_args(FLAGS, sys.argv)
  if FLAGS.model_type == 'variational_dropout':
    fs = VariationalDropoutFS(
        FLAGS.config_path,
        FLAGS.output_dir,
        FLAGS.topk,
        checkpoint_path=FLAGS.checkpoint_path,
        fg_path=FLAGS.fg_path,
        visualize=FLAGS.visualize)
    fs.process()
  else:
    raise ValueError('Unknown feature selection model type %s' %
                     FLAGS.model_type)