From 00eccfdb0d315de3d94bfa1da84cc1dc65c8af39 Mon Sep 17 00:00:00 2001 From: Ryan Curtin Date: Tue, 20 Oct 2015 13:41:35 +0000 Subject: [PATCH] Refactor allkfn program. --- .../methods/neighbor_search/allkfn_main.cpp | 327 +++++++++--------- 1 file changed, 167 insertions(+), 160 deletions(-) diff --git a/src/mlpack/methods/neighbor_search/allkfn_main.cpp b/src/mlpack/methods/neighbor_search/allkfn_main.cpp index 784faef467..5bf298213e 100644 --- a/src/mlpack/methods/neighbor_search/allkfn_main.cpp +++ b/src/mlpack/methods/neighbor_search/allkfn_main.cpp @@ -13,6 +13,7 @@ #include "neighbor_search.hpp" #include "unmap.hpp" +#include "ns_model.hpp" using namespace std; using namespace mlpack; @@ -41,199 +42,205 @@ PROGRAM_INFO("All K-Furthest-Neighbors", "corresponds to the distance between those two points."); // Define our input parameters that this program will take. -PARAM_STRING_REQ("reference_file", "File containing the reference dataset.", - "r"); -PARAM_INT_REQ("k", "Number of furthest neighbors to find.", "k"); -PARAM_STRING_REQ("distances_file", "File to output distances into.", "d"); -PARAM_STRING_REQ("neighbors_file", "File to output neighbors into.", "n"); +PARAM_STRING("reference_file", "File containing the reference dataset.", "r", + ""); +PARAM_STRING("distances_file", "File to output distances into.", "d", ""); +PARAM_STRING("neighbors_file", "File to output neighbors into.", "n", ""); +// The option exists to load or save models. +PARAM_STRING("input_model_file", "File containing pre-trained kFN model.", "m", + ""); +PARAM_STRING("output_model_file", "If specified, the kFN model will be saved to" + " the given file.", "M", ""); + +// The user may specify a query file of query points and a number of furthest +// neighbors to search for. PARAM_STRING("query_file", "File containing query points (optional).", "q", ""); +PARAM_INT("k", "Number of furthest neighbors to find.", "k", 0); +// The user may specify the type of tree to use, and a few pararmeters for tree +// building. +PARAM_STRING("tree_type", "Type of tree to use: 'kd', 'cover', 'r', 'r-star', " + "'ball'.", "t", "kd"); PARAM_INT("leaf_size", "Leaf size for tree building.", "l", 20); +PARAM_FLAG("random_basis", "Before tree-building, project the data onto a " + "random orthogonal basis.", "R"); +PARAM_INT("seed", "Random seed (if 0, std::time(NULL) is used).", "s", 0); + +// Search settings. PARAM_FLAG("naive", "If true, O(n^2) naive mode is used for computation.", "N"); PARAM_FLAG("single_mode", "If true, single-tree search is used (as opposed to " "dual-tree search).", "s"); -PARAM_FLAG("r_tree", "If true, use an R-Tree to perform the search " - "(experimental, may be slow.).", "T"); + +// Convenience typedef. +typedef NSModel KFNModel; int main(int argc, char *argv[]) { // Give CLI the command line parameters the user passed in. CLI::ParseCommandLine(argc, argv); - // Get all the parameters. - string referenceFile = CLI::GetParam("reference_file"); + if (CLI::GetParam("seed") != 0) + math::RandomSeed((size_t) CLI::GetParam("seed")); + else + math::RandomSeed((size_t) std::time(NULL)); - string distancesFile = CLI::GetParam("distances_file"); - string neighborsFile = CLI::GetParam("neighbors_file"); + // A user cannot specify both reference data and a model. + if (CLI::HasParam("reference_file") && CLI::HasParam("input_model_file")) + Log::Fatal << "Only one of --reference_file (-r) or --input_model_file (-m)" + << " may be specified!" << endl; - int lsInt = CLI::GetParam("leaf_size"); + // A user must specify one of them... + if (!CLI::HasParam("reference_file") && !CLI::HasParam("input_model_file")) + Log::Fatal << "No model specified (--input_model_file) and no reference " + << "data specified (--reference_file)! One must be provided." << endl; - size_t k = CLI::GetParam("k"); - - bool naive = CLI::HasParam("naive"); - bool singleMode = CLI::HasParam("single_mode"); - - arma::mat referenceData; - arma::mat queryData; // So it doesn't go out of scope. - data::Load(referenceFile, referenceData, true); - - Log::Info << "Loaded reference data from '" << referenceFile << "' (" - << referenceData.n_rows << " x " << referenceData.n_cols << ")." << endl; - - // Sanity check on k value: must be greater than 0, must be less than the - // number of reference points. - if (k > referenceData.n_cols) + if (CLI::HasParam("input_model_file")) { - Log::Fatal << "Invalid k: " << k << "; must be greater than 0 and less "; - Log::Fatal << "than or equal to the number of reference points ("; - Log::Fatal << referenceData.n_cols << ")." << endl; + // Notify the user of parameters that will be ignored. + if (CLI::HasParam("tree_type")) + Log::Warn << "--tree_type (-t) will be ignored because --input_model_file" + << " is specified." << endl; + if (CLI::HasParam("leaf_size")) + Log::Warn << "--leaf_size (-l) will be ignored because --input_model_file" + << " is specified." << endl; + if (CLI::HasParam("random_basis")) + Log::Warn << "--random_basis (-R) will be ignored because " + << "--input_model_file is specified." << endl; + if (CLI::HasParam("naive")) + Log::Warn << "--naive (-N) will be ignored because --input_model_file is " + << "specified." << endl; } - if (CLI::GetParam("query_file") != "") - { - string queryFile = CLI::GetParam("query_file"); - data::Load(queryFile, queryData, true); - } + // The user should give something to do... + if (!CLI::HasParam("k") && !CLI::HasParam("output_model_file")) + Log::Warn << "Neither -k nor --output_model_file are specified, so no " + << "results from this program will be saved!" << endl; + + // If the user specifies k but no output files, they should be warned. + if (CLI::HasParam("k") && + !(CLI::HasParam("neighbors_file") || CLI::HasParam("distances_file"))) + Log::Warn << "Neither --neighbors_file nor --distances_file is specified, " + << "so the furthest neighbor search results will not be saved!" << endl; + + // If the user specifies output files but no k, they should be warned. + if ((CLI::HasParam("neighbors_file") || CLI::HasParam("distances_file")) && + !CLI::HasParam("k")) + Log::Warn << "An output file for furthest neighbor search is given (" + << "--neighbors_file or --distances_file), but furthest neighbor search" + << " is not being performed because k (--k) is not specified!" << endl; // Sanity check on leaf size. - if (lsInt < 0) + const int lsInt = CLI::GetParam("leaf_size"); + if (lsInt < 1) + Log::Fatal << "Invalid leaf size: " << lsInt << ". Must be greater than 0." + << endl; + + // We either have to load the reference data, or we have to load the model. + NSModel kfn; + const bool naive = CLI::HasParam("naive"); + const bool singleMode = CLI::HasParam("single_mode"); + if (CLI::HasParam("reference_file")) { - Log::Fatal << "Invalid leaf size: " << lsInt << ". Must be greater " - "than or equal to 0." << endl; - } - size_t leafSize = lsInt; + // Get all the parameters. + const string referenceFile = CLI::GetParam("reference_file"); + const string treeType = CLI::GetParam("tree_type"); + const bool randomBasis = CLI::HasParam("random_basis"); - // Naive mode overrides single mode. - if (singleMode && naive) - { - Log::Warn << "--single_mode ignored because --naive is present." << endl; - } - - arma::Mat neighbors; - arma::mat distances; - - if (naive) - { - AllkFN allkfn(referenceData, false, naive); - - if (CLI::HasParam("query_file")) - allkfn.Search(queryData, k, neighbors, distances); + int tree = 0; + if (treeType == "kd") + tree = KFNModel::KD_TREE; + else if (treeType == "cover") + tree = KFNModel::COVER_TREE; + else if (treeType == "r") + tree = KFNModel::R_TREE; + else if (treeType == "r-star") + tree = KFNModel::R_STAR_TREE; + else if (treeType == "ball") + tree = KFNModel::BALL_TREE; else - allkfn.Search(k, neighbors, distances); - } - if (!CLI::HasParam("r_tree")) - { - // Use default kd-tree. - std::vector oldFromNewRefs; + Log::Fatal << "Unknown tree type '" << treeType << "'; valid choices are " + << "'kd', 'cover', 'r', 'r-star', and 'ball'." << endl; - typedef KDTree, - arma::mat> TreeType; + kfn.TreeType() = tree; + kfn.RandomBasis() = randomBasis; - // Build trees by hand, so we can save memory: if we pass a tree to - // NeighborSearch, it does not copy the matrix. - Log::Info << "Building reference tree..." << endl; - Timer::Start("reference_tree_building"); - TreeType refTree(referenceData, oldFromNewRefs, leafSize); - Timer::Stop("reference_tree_building"); + arma::mat referenceSet; + data::Load(referenceFile, referenceSet, true); - std::vector oldFromNewQueries; + Log::Info << "Loaded reference data from '" << referenceFile << "' (" + << referenceSet.n_rows << " x " << referenceSet.n_cols << ")." << endl; - AllkFN allkfn(&refTree, singleMode); - - arma::mat distancesOut(distances.n_rows, distances.n_cols); - arma::Mat neighborsOut(neighbors.n_rows, neighbors.n_cols); - - if (CLI::HasParam("query_file")) - { - if (!singleMode) - { - // Build trees by hand, so we can save memory: if we pass a tree to - // NeighborSearch, it does not copy the matrix. - Log::Info << "Building query tree..." << endl; - Timer::Start("tree_building"); - TreeType queryTree(queryData, oldFromNewQueries, leafSize); - Timer::Stop("tree_building"); - Log::Info << "Tree built." << endl; - - Log::Info << "Computing " << k << " furthest neighbors..." << endl; - allkfn.Search(&queryTree, k, neighborsOut, distancesOut); - } - else - { - Log::Info << "Computing " << k << " furthest neighbors..." << endl; - allkfn.Search(queryData, k, neighborsOut, distancesOut); - } - } - else - { - Log::Info << "Computing " << k << " furthest neighbors..." << endl; - allkfn.Search(k, neighborsOut, distancesOut); - } - - Log::Info << "Neighbors computed." << endl; - - // We have to map back to the original indices from before the tree - // construction. - Log::Info << "Re-mapping indices..." << endl; - - // Map the points back to their original locations. - if ((CLI::GetParam("query_file") != "") && !singleMode) - Unmap(neighborsOut, distancesOut, oldFromNewRefs, oldFromNewQueries, neighbors, - distances); - else if ((CLI::GetParam("query_file") != "") && singleMode) - Unmap(neighborsOut, distancesOut, oldFromNewRefs, neighbors, distances); - else - Unmap(neighborsOut, distancesOut, oldFromNewRefs, oldFromNewRefs, neighbors, - distances); + const size_t leafSize = (size_t) lsInt; + kfn.BuildModel(std::move(referenceSet), leafSize, naive, singleMode); } else { - // Use the R tree. - Log::Info << "Using R tree for furthest-neighbor calculation." << endl; + // Load the model from file. + const string inputModelFile = CLI::GetParam("input_model_file"); + data::Load(inputModelFile, "kfn_model", kfn, true); // Fatal on failure. - // Convenience typedef. - typedef RStarTree, arma::mat> TreeType; + Log::Info << "Loaded kFN model from '" << inputModelFile << "' (trained on " + << kfn.Dataset().n_rows << "x" << kfn.Dataset().n_cols << " dataset)." + << endl; - // Build trees by hand, so we can save memory: if we pass a tree to - // NeighborSearch, it does not copy the matrix. - Log::Info << "Building reference tree..." << endl; - Timer::Start("tree_building"); - TreeType refTree(referenceData, leafSize, leafSize * 0.4, 5, 2, 0); - Timer::Stop("tree_building"); - Log::Info << "Tree built." << endl; - - typedef NeighborSearch AllkFNType; - AllkFNType allkfn(&refTree, singleMode); - - if (CLI::GetParam("query_file") != "") - { - if (!singleMode) - { - Timer::Start("tree_building"); - TreeType queryTree(queryData, leafSize, leafSize * 0.4, 5, 2, 0); - Timer::Stop("tree_building"); - - Log::Info << "Computing " << k << " nearest neighbors..." << endl; - allkfn.Search(&queryTree, k, neighbors, distances); - } - else - { - Log::Info << "Computing " << k << " nearest neighbors..." << endl; - allkfn.Search(queryData, k, neighbors, distances); - } - } - else - { - Log::Info << "Computing " << k << " nearest neighbors..." << endl; - allkfn.Search(k, neighbors, distances); - } - Log::Info << "Neighbors computed." << endl; + // Adjust singleMode and naive if necessary. + if (CLI::HasParam("single_mode")) + kfn.SingleMode() = true; + if (CLI::HasParam("naive")) + kfn.Naive() = true; + if (CLI::HasParam("leaf_size")) + kfn.LeafSize() = (size_t) lsInt; } - // Save output. - data::Save(distancesFile, distances); - data::Save(neighborsFile, neighbors); + // Perform search, if desired. + if (CLI::HasParam("k")) + { + const string queryFile = CLI::GetParam("query_file"); + const size_t k = (size_t) CLI::GetParam("k"); + + arma::mat queryData; + if (queryFile != "") + { + data::Load(queryFile, queryData, true); + Log::Info << "Loaded query data from '" << queryFile << "' (" + << queryData.n_rows << " x " << queryData.n_cols << ")." << endl; + } + + // Sanity check on k value: must be greater than 0, must be less than the + // number of reference points. Since it is unsigned, we only test the upper + // bound. + if (k > kfn.Dataset().n_cols) + { + Log::Fatal << "Invalid k: " << k << "; must be greater than 0 and less " + << "than or equal to the number of reference points (" + << kfn.Dataset().n_cols << ")." << endl; + } + + // Naive mode overrides single mode. + if (singleMode && naive) + Log::Warn << "--single_mode ignored because --naive is present." << endl; + + // Now run the search. + arma::Mat neighbors; + arma::mat distances; + + if (CLI::HasParam("query_file")) + kfn.Search(std::move(queryData), k, neighbors, distances); + else + kfn.Search(k, neighbors, distances); + Log::Info << "Search complete." << endl; + + // Save output, if desired. + if (CLI::HasParam("neighbors_file")) + data::Save(CLI::GetParam("neighbors_file"), neighbors); + if (CLI::HasParam("distances_file")) + data::Save(CLI::GetParam("distances_file"), distances); + } + + if (CLI::HasParam("output_model_file")) + { + const string outputModelFile = CLI::GetParam("output_model_File"); + data::Save(outputModelFile, "kfn_model", kfn); + } }