Refactor allkfn program.
This commit is contained in:
Ryan Curtin
@@ -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<FurthestNeighborSort> 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<string>("reference_file");
|
||||
if (CLI::GetParam<int>("seed") != 0)
|
||||
math::RandomSeed((size_t) CLI::GetParam<int>("seed"));
|
||||
else
|
||||
math::RandomSeed((size_t) std::time(NULL));
|
||||
|
||||
string distancesFile = CLI::GetParam<string>("distances_file");
|
||||
string neighborsFile = CLI::GetParam<string>("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<int>("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<int>("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<string>("query_file") != "")
|
||||
{
|
||||
string queryFile = CLI::GetParam<string>("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<int>("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<FurthestNeighborSort> 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<string>("reference_file");
|
||||
const string treeType = CLI::GetParam<string>("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<size_t> 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<size_t> oldFromNewRefs;
|
||||
Log::Fatal << "Unknown tree type '" << treeType << "'; valid choices are "
|
||||
<< "'kd', 'cover', 'r', 'r-star', and 'ball'." << endl;
|
||||
|
||||
typedef KDTree<EuclideanDistance, NeighborSearchStat<FurthestNeighborSort>,
|
||||
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<size_t> 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<size_t> 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<string>("query_file") != "") && !singleMode)
|
||||
Unmap(neighborsOut, distancesOut, oldFromNewRefs, oldFromNewQueries, neighbors,
|
||||
distances);
|
||||
else if ((CLI::GetParam<string>("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<string>("input_model_file");
|
||||
data::Load(inputModelFile, "kfn_model", kfn, true); // Fatal on failure.
|
||||
|
||||
// Convenience typedef.
|
||||
typedef RStarTree<EuclideanDistance,
|
||||
NeighborSearchStat<FurthestNeighborSort>, 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<FurthestNeighborSort, EuclideanDistance, arma::mat,
|
||||
RStarTree> AllkFNType;
|
||||
AllkFNType allkfn(&refTree, singleMode);
|
||||
|
||||
if (CLI::GetParam<string>("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<string>("query_file");
|
||||
const size_t k = (size_t) CLI::GetParam<int>("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<size_t> 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<string>("neighbors_file"), neighbors);
|
||||
if (CLI::HasParam("distances_file"))
|
||||
data::Save(CLI::GetParam<string>("distances_file"), distances);
|
||||
}
|
||||
|
||||
if (CLI::HasParam("output_model_file"))
|
||||
{
|
||||
const string outputModelFile = CLI::GetParam<string>("output_model_File");
|
||||
data::Save(outputModelFile, "kfn_model", kfn);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user