Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Autocast numpy.ndarray <-> std::vector<Eigen::Vector> #248

Merged
merged 3 commits into from
Jan 24, 2024

Conversation

sarlinpe
Copy link
Collaborator

@sarlinpe sarlinpe commented Jan 24, 2024

For large inputs, the default pybind type caster is much slower than directly copying the underlying memory segment, as hinted by pybind/pybind11#1481.

Test code:

  using Matrix =
      typename Eigen::Matrix<double, Eigen::Dynamic, 2, Eigen::RowMajor>;
  using Vector = typename Eigen::Vector2d;
  m.def("test_vector", [](const std::vector<Vector>& x) { return x; });
  m.def(
      "test_map",
      [](const Eigen::Ref<Matrix>& x) {
        // Cast to std::vector
        std::vector<Vector> vec(x.rows());
        Eigen::Map<Matrix>(
            reinterpret_cast<double*>(vec.data()), vec.size(), 2) = x;
        // Cast back to Eigen::Matrix
        return Eigen::Map<const Matrix>(
            reinterpret_cast<const double*>(vec.data()), vec.size(), 2);
      },
      py::return_value_policy::copy);
  m.def("test_matrix", [](const Eigen::Ref<Matrix>& x) { return x; });
n = 100_000
print(f'{"size":>5}: vector   map      matrix')
for i in [10, 100, 1_000, 10_000]:
    x = np.random.randn(i, 2)
    assert np.all(pycolmap.test_vector(x) == x)
    t_vec = timeit.timeit(lambda: pycolmap.test_vector(x), number=n)/n
    t_map = timeit.timeit(lambda: pycolmap.test_map(x), number=n)/n
    t_mat = timeit.timeit(lambda: pycolmap.test_matrix(x), number=n)/n
    print(f'{i:>5}: {t_vec:.2E} {t_map:.2E} {t_mat:.2E}')

Timings (second) without the type caster:

 size: vector   map      matrix
   10: 6.95E-06 7.71E-07 4.90E-07
  100: 6.90E-05 7.15E-07 5.28E-07
 1000: 6.82E-04 1.04E-06 5.16E-07
10000: 7.04E-03 5.77E-06 5.03E-07

With:

 size: vector
   10: 1.02E-06
  100: 1.29E-06
 1000: 2.25E-06
10000: 8.94E-05

This is a 100x speedup for 10k inputs. There is still a gap w.r.t direct mapping for C++ -> Python, but I haven't managed to find the source. Ideally COLMAP would use Eigen::Matrix when possible.

@sarlinpe sarlinpe marked this pull request as ready for review January 24, 2024 15:47
@sarlinpe sarlinpe merged commit 27e36eb into master Jan 24, 2024
5 checks passed
@sarlinpe sarlinpe deleted the sarlinpe/autocast-vec-vec branch January 24, 2024 16:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants