diff --git a/src/mlpack/methods/fastmks/fastmks_main.cpp b/src/mlpack/methods/fastmks/fastmks_main.cpp index ada3eb7852..73d81bdc95 100644 --- a/src/mlpack/methods/fastmks/fastmks_main.cpp +++ b/src/mlpack/methods/fastmks/fastmks_main.cpp @@ -7,6 +7,7 @@ #include #include "fastmks.hpp" +#include "fastmks_model.hpp" using namespace std; using namespace mlpack; @@ -39,25 +40,11 @@ PROGRAM_INFO("FastMKS (Fast Max-Kernel Search)", "This executable performs FastMKS using a cover tree. The base used to " "build the cover tree can be specified with the --base option."); -// Define our input parameters. -PARAM_STRING_REQ("reference_file", "File containing the reference dataset.", - "r"); -PARAM_STRING("query_file", "File containing the query dataset.", "q", ""); - -PARAM_INT_REQ("k", "Number of maximum kernels to find.", "k"); - -PARAM_STRING("kernels_file", "File to save kernels into.", "p", ""); -PARAM_STRING("indices_file", "File to save indices of kernels into.", - "i", ""); - +// Model-building parameters. +PARAM_STRING("reference_file", "File containing the reference dataset.", "r", + ""); PARAM_STRING("kernel", "Kernel type to use: 'linear', 'polynomial', 'cosine', " "'gaussian', 'epanechnikov', 'triangular', 'hyptan'.", "K", "linear"); - -PARAM_FLAG("naive", "If true, O(n^2) naive mode is used for computation.", "N"); -PARAM_FLAG("single", "If true, single-tree search is used (as opposed to " - "dual-tree search.", "S"); - -// Cover tree parameter. PARAM_DOUBLE("base", "Base to use during cover tree construction.", "b", 2.0); // Kernel parameters. @@ -68,262 +55,197 @@ PARAM_DOUBLE("bandwidth", "Bandwidth (for Gaussian, Epanechnikov, and " "triangular kernels).", "w", 1.0); PARAM_DOUBLE("scale", "Scale of kernel (for hyptan kernel).", "s", 1.0); -//! Run FastMKS on a single dataset for the given kernel type. -template -void RunFastMKS(const arma::mat& referenceData, - const bool single, - const bool naive, - const double base, - const size_t k, - arma::Mat& indices, - arma::mat& kernels, - KernelType& kernel) -{ - if (naive) - { - // No need for trees. - FastMKS fastmks(referenceData, kernel, false, naive); - fastmks.Search(k, indices, kernels); - } - else - { - // Create the tree with the specified base. - typedef CoverTree, FastMKSStat, arma::mat, - FirstPointIsRoot> TreeType; - IPMetric metric(kernel); - TreeType tree(referenceData, metric, base); +// Load/save models. +PARAM_STRING("input_model_file", "File containing FastMKS model.", "m", ""); +PARAM_STRING("output_model_file", "File to save FastMKS model to.", "M", ""); - // Create FastMKS object. - FastMKS fastmks(&tree, single); +// Search preferences. +PARAM_STRING("query_file", "File containing the query dataset.", "q", ""); +PARAM_INT("k", "Number of maximum kernels to find.", "k", 0); +PARAM_FLAG("naive", "If true, O(n^2) naive mode is used for computation.", "N"); +PARAM_FLAG("single", "If true, single-tree search is used (as opposed to " + "dual-tree search.", "S"); - // Now search with it. - fastmks.Search(k, indices, kernels); - } -} - -//! Run FastMKS for a given query and reference set using the given kernel type. -template -void RunFastMKS(const arma::mat& referenceData, - const arma::mat& queryData, - const bool single, - const bool naive, - const double base, - const size_t k, - arma::Mat& indices, - arma::mat& kernels, - KernelType& kernel) -{ - if (naive) - { - // No need for trees. - FastMKS fastmks(referenceData, kernel, false, naive); - fastmks.Search(queryData, k, indices, kernels); - } - else - { - // Create the tree with the specified base. - typedef CoverTree, FastMKSStat, arma::mat, - FirstPointIsRoot> TreeType; - IPMetric metric(kernel); - TreeType referenceTree(referenceData, metric, base); - - // Create FastMKS object. - FastMKS fastmks(&referenceTree, - single); - - // Now search with it. - if (single) - { - fastmks.Search(queryData, k, indices, kernels); - } - else - { - TreeType queryTree(queryData, metric, base); - fastmks.Search(&queryTree, k, indices, kernels); - } - } -} +PARAM_STRING("kernels_file", "File to save kernels into.", "p", ""); +PARAM_STRING("indices_file", "File to save indices of kernels into.", + "i", ""); int main(int argc, char** argv) { CLI::ParseCommandLine(argc, argv); - // Get reference dataset filename. - const string referenceFile = CLI::GetParam("reference_file"); + // Validate command-line parameters. + if (CLI::HasParam("reference_file") && CLI::HasParam("input_model_file")) + Log::Fatal << "Cannot specify both --reference_file (-r) and " + << "--input_model_file (-m)!" << endl; - // The number of max kernel values to find. - const size_t k = CLI::GetParam("k"); + if (!CLI::HasParam("reference_file") && !CLI::HasParam("input_model_file")) + Log::Fatal << "Must specify either --reference_file (-r) or " + << "--input_model_file (-m)!" << endl; - // Runtime parameters. - const bool naive = CLI::HasParam("naive"); - const bool single = CLI::HasParam("single"); - - // For cover tree construction. - const double base = CLI::GetParam("base"); - - // Kernel parameters. - const string kernelType = CLI::GetParam("kernel"); - const double degree = CLI::GetParam("degree"); - const double offset = CLI::GetParam("offset"); - const double bandwidth = CLI::GetParam("bandwidth"); - const double scale = CLI::GetParam("scale"); - - // The datasets. The query matrix may never be used. - arma::mat referenceData; - arma::mat queryData; - - 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. - 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; + if (CLI::HasParam("kernel")) + Log::Warn << "--kernel (-k) ignored because --input_model_file (-m) is " + << "specified." << endl; + if (CLI::HasParam("bandwidth")) + Log::Warn << "--bandwidth (-w) ignored because --input_model_file (-m) is" + << " specified." << endl; + if (CLI::HasParam("degree")) + Log::Warn << "--degree (-d) ignored because --input_model_file (-m) is " + << " specified." << endl; + if (CLI::HasParam("offset")) + Log::Warn << "--offset (-o) ignored because --input_model_file (-m) is " + << " specified." << endl; } + if (!CLI::HasParam("k") && + (CLI::HasParam("indices_file") || CLI::HasParam("kernels_file"))) + Log::Warn << "--indices_file and --kernels_file ignored, because no search " + << "task is specified (i.e., --k is not specified)!" << endl; + + if (CLI::HasParam("k") && + !(CLI::HasParam("indices_file") || CLI::HasParam("kernels_file"))) + Log::Warn << "Search specified with --k, but no output will be saved " + << "because neither --indices_file nor --kernels_file are specified!" + << endl; + // Check on kernel type. + const string kernelType = CLI::GetParam("kernel"); if ((kernelType != "linear") && (kernelType != "polynomial") && (kernelType != "cosine") && (kernelType != "gaussian") && - (kernelType != "graph") && (kernelType != "approxGraph") && (kernelType != "triangular") && (kernelType != "hyptan") && - (kernelType != "inv-mq") && (kernelType != "epanechnikov")) + (kernelType != "epanechnikov")) { - Log::Fatal << "Invalid kernel type: '" << kernelType << "'; must be "; - Log::Fatal << "'linear' or 'polynomial'." << endl; - } - - // Load the query matrix, if we can. - if (CLI::HasParam("query_file")) - { - const string queryFile = CLI::GetParam("query_file"); - data::Load(queryFile, queryData, true); - - Log::Info << "Loaded query data from '" << queryFile << "' (" - << queryData.n_rows << " x " << queryData.n_cols << ")." << endl; - } - else - { - Log::Info << "Using reference dataset as query dataset (--query_file not " - << "specified)." << endl; + Log::Fatal << "Invalid kernel type: '" << kernelType << "'; must be " + << "'linear', 'polynomial', 'cosine', 'gaussian', 'triangular', or " + << "'epanechnikov'." << endl; } // Naive mode overrides single mode. - if (naive && single) - { + if (CLI::HasParam("naive") && CLI::HasParam("single")) Log::Warn << "--single ignored because --naive is present." << endl; - } - // Matrices for output storage. - arma::Mat indices; - arma::mat kernels; - - // Construct FastMKS object. - if (queryData.n_elem == 0) + FastMKSModel model; + arma::mat referenceData; + if (CLI::HasParam("reference_file")) { + data::Load(CLI::GetParam("reference_file"), referenceData, true); + + Log::Info << "Loaded reference data from '" + << CLI::GetParam("reference_file") << "' (" + << referenceData.n_rows << " x " << referenceData.n_cols << ")." + << endl; + + // For cover tree construction. + const double base = CLI::GetParam("base"); + + // Kernel parameters. + const string kernelType = CLI::GetParam("kernel"); + const double degree = CLI::GetParam("degree"); + const double offset = CLI::GetParam("offset"); + const double bandwidth = CLI::GetParam("bandwidth"); + const double scale = CLI::GetParam("scale"); + + // Search preferences. + const bool naive = CLI::HasParam("naive"); + const bool single = CLI::HasParam("single"); + if (kernelType == "linear") { LinearKernel lk; - RunFastMKS(referenceData, single, naive, base, k, indices, - kernels, lk); + model.KernelType() = FastMKSModel::LINEAR_KERNEL; + model.BuildModel(referenceData, lk, single, naive, base); } else if (kernelType == "polynomial") { - PolynomialKernel pk(degree, offset); - RunFastMKS(referenceData, single, naive, base, k, - indices, kernels, pk); + model.KernelType() = FastMKSModel::POLYNOMIAL_KERNEL; + model.BuildModel(referenceData, pk, single, naive, base); } else if (kernelType == "cosine") { CosineDistance cd; - RunFastMKS(referenceData, single, naive, base, k, indices, - kernels, cd); + model.KernelType() = FastMKSModel::COSINE_DISTANCE; + model.BuildModel(referenceData, cd, single, naive, base); } else if (kernelType == "gaussian") { GaussianKernel gk(bandwidth); - RunFastMKS(referenceData, single, naive, base, k, indices, - kernels, gk); + model.KernelType() = FastMKSModel::GAUSSIAN_KERNEL; + model.BuildModel(referenceData, gk, single, naive, base); } else if (kernelType == "epanechnikov") { EpanechnikovKernel ek(bandwidth); - RunFastMKS(referenceData, single, naive, base, k, - indices, kernels, ek); + model.KernelType() = FastMKSModel::EPANECHNIKOV_KERNEL; + model.BuildModel(referenceData, ek, single, naive, base); } else if (kernelType == "triangular") { TriangularKernel tk(bandwidth); - RunFastMKS(referenceData, single, naive, base, k, - indices, kernels, tk); + model.KernelType() = FastMKSModel::TRIANGULAR_KERNEL; + model.BuildModel(referenceData, tk, single, naive, base); } else if (kernelType == "hyptan") { HyperbolicTangentKernel htk(scale, offset); - RunFastMKS(referenceData, single, naive, base, k, - indices, kernels, htk); + model.KernelType() = FastMKSModel::HYPTAN_KERNEL; + model.BuildModel(referenceData, htk, single, naive, base); } } else { - if (kernelType == "linear") + // Load model from file, then do whatever is necessary. + data::Load(CLI::GetParam("input_model_file"), "fastmks_model", + model, true); + } + + // Set search preferences. + model.Naive() = CLI::HasParam("naive"); + model.SingleMode() = CLI::HasParam("single"); + + // Should we do search? + if (CLI::HasParam("k")) + { + arma::mat kernels; + arma::Mat indices; + + if (CLI::HasParam("query_file")) { - LinearKernel lk; - RunFastMKS(referenceData, queryData, single, naive, base, k, - indices, kernels, lk); + const string queryFile = CLI::GetParam("query_file"); + const double base = CLI::GetParam("base"); + + arma::mat queryData; + data::Load(queryFile, queryData, true); + + Log::Info << "Loaded query data from '" << queryFile << "' (" + << queryData.n_rows << " x " << queryData.n_cols << ")." << endl; + + model.Search(queryData, (size_t) CLI::GetParam("k"), indices, + kernels, base); } - else if (kernelType == "polynomial") + else { - PolynomialKernel pk(degree, offset); - RunFastMKS(referenceData, queryData, single, naive, - base, k, indices, kernels, pk); + model.Search((size_t) CLI::GetParam("k"), indices, kernels); } - else if (kernelType == "cosine") + + // Save output, if we were asked to. + if (CLI::HasParam("kernels_file")) { - CosineDistance cd; - RunFastMKS(referenceData, queryData, single, naive, base, - k, indices, kernels, cd); + const string kernelsFile = CLI::GetParam("kernels_file"); + data::Save(kernelsFile, kernels, false); } - else if (kernelType == "gaussian") + + if (CLI::HasParam("indices_file")) { - GaussianKernel gk(bandwidth); - RunFastMKS(referenceData, queryData, single, naive, base, - k, indices, kernels, gk); - } - else if (kernelType == "epanechnikov") - { - EpanechnikovKernel ek(bandwidth); - RunFastMKS(referenceData, queryData, single, naive, - base, k, indices, kernels, ek); - } - else if (kernelType == "triangular") - { - TriangularKernel tk(bandwidth); - RunFastMKS(referenceData, queryData, single, naive, - base, k, indices, kernels, tk); - } - else if (kernelType == "hyptan") - { - HyperbolicTangentKernel htk(scale, offset); - RunFastMKS(referenceData, queryData, single, - naive, base, k, indices, kernels, htk); + const string indicesFile = CLI::GetParam("indices_file"); + data::Save(indicesFile, indices, false); } } - // Save output, if we were asked to. - if (CLI::HasParam("kernels_file")) - { - const string kernelsFile = CLI::GetParam("kernels_file"); - data::Save(kernelsFile, kernels, false); - } - - if (CLI::HasParam("indices_file")) - { - const string indicesFile = CLI::GetParam("indices_file"); - data::Save(indicesFile, indices, false); - } + // Save the model, if requested. + if (CLI::HasParam("output_model_file")) + data::Save(CLI::GetParam("output_model_file"), "fastmks_model", + model); }