hoeffding_tree_model.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_HOEFFDING_TREE_HOEFFDING_TREE_MODEL_HPP
#define MLPACK_METHODS_HOEFFDING_TREE_HOEFFDING_TREE_MODEL_HPP

#include "hoeffding_tree.hpp"
#include "binary_numeric_split.hpp"
#include "information_gain.hpp"

namespace mlpack {
namespace tree {

class HoeffdingTreeModel
{
 public:
  enum TreeType
  {
    GINI_HOEFFDING,
    GINI_BINARY,
    INFO_HOEFFDING,
    INFO_BINARY
  };

  typedef HoeffdingTree<GiniImpurity, HoeffdingDoubleNumericSplit,
      HoeffdingCategoricalSplit> GiniHoeffdingTreeType;
  typedef HoeffdingTree<GiniImpurity, BinaryDoubleNumericSplit,
      HoeffdingCategoricalSplit> GiniBinaryTreeType;
  typedef HoeffdingTree<InformationGain, HoeffdingDoubleNumericSplit,
      HoeffdingCategoricalSplit> InfoHoeffdingTreeType;
  typedef HoeffdingTree<InformationGain, BinaryDoubleNumericSplit,
      HoeffdingCategoricalSplit> InfoBinaryTreeType;

  HoeffdingTreeModel(const TreeType& type = GINI_HOEFFDING);

  HoeffdingTreeModel(const HoeffdingTreeModel& other);

  HoeffdingTreeModel(HoeffdingTreeModel&& other);

  HoeffdingTreeModel& operator=(const HoeffdingTreeModel& other);

  HoeffdingTreeModel& operator=(HoeffdingTreeModel&& other);

  ~HoeffdingTreeModel();

  void BuildModel(const arma::mat& dataset,
                  const data::DatasetInfo& datasetInfo,
                  const arma::Row<size_t>& labels,
                  const size_t numClasses,
                  const bool batchTraining,
                  const double successProbability,
                  const size_t maxSamples,
                  const size_t checkInterval,
                  const size_t minSamples,
                  const size_t bins,
                  const size_t observationsBeforeBinning);

  void Train(const arma::mat& dataset,
             const arma::Row<size_t>& labels,
             const bool batchTraining);

  void Classify(const arma::mat& dataset,
                arma::Row<size_t>& predictions) const;

  void Classify(const arma::mat& dataset,
                arma::Row<size_t>& predictions,
                arma::rowvec& probabilities) const;

  size_t NumNodes() const;

  template<typename Archive>
  void Serialize(Archive& ar, const unsigned int /* version */)
  {
    ar & data::CreateNVP(type, "type");

    // Clear memory if needed.
    if (Archive::is_loading::value)
    {
      delete giniHoeffdingTree;
      delete giniBinaryTree;
      delete infoHoeffdingTree;
      delete infoBinaryTree;

      giniHoeffdingTree = NULL;
      giniBinaryTree = NULL;
      infoHoeffdingTree = NULL;
      infoBinaryTree = NULL;
    }

    // Fake dataset info may be needed to create fake trees.
    data::DatasetInfo info;
    if (type == GINI_HOEFFDING)
    {
      // Create fake tree to load into if needed.
      if (Archive::is_loading::value)
        giniHoeffdingTree = new GiniHoeffdingTreeType(info, 1, 1);
      ar & data::CreateNVP(*giniHoeffdingTree, "giniHoeffdingTree");
    }
    else if (type == GINI_BINARY)
    {
      // Create fake tree to load into if needed.
      if (Archive::is_loading::value)
        giniBinaryTree = new GiniBinaryTreeType(info, 1, 1);
      ar & data::CreateNVP(*giniBinaryTree, "giniBinaryTree");
    }
    else if (type == INFO_HOEFFDING)
    {
      // Create fake tree to load into if needed.
      if (Archive::is_loading::value)
        infoHoeffdingTree = new InfoHoeffdingTreeType(info, 1, 1);
      ar & data::CreateNVP(*infoHoeffdingTree, "infoHoeffdingTree");
    }
    else if (type == INFO_BINARY)
    {
      // Create fake tree to load into if needed.
      if (Archive::is_loading::value)
        infoBinaryTree = new InfoBinaryTreeType(info, 1, 1);
      ar & data::CreateNVP(*infoBinaryTree, "infoBinaryTree");
    }
  }

 private:
  TreeType type;

  GiniHoeffdingTreeType* giniHoeffdingTree;

  GiniBinaryTreeType* giniBinaryTree;

  InfoHoeffdingTreeType* infoHoeffdingTree;

  InfoBinaryTreeType* infoBinaryTree;
};

} // namespace tree
} // namespace mlpack

#endif