Very slow torch.median() compared to CuPy

crcrpar · March 17, 2021, 6:42am

Disclaimer: I didn’t thoroughly read the source code, I just git grep-ed and googled some information. So what I say might be wrong.

PyTorchv1.8’s median calls sort internally:

github.com

pytorch/pytorch/blob/37c1f4a7fef115d719104e871d0cf39434aa9d56/aten/src/ATen/native/cuda/Sorting.cu#L393


  return std::forward_as_tuple(values, indices);
}

Tensor median_impl(const Tensor& self, bool ignore_nan) {
  NoNamesGuard guard;

  int64_t size = self.numel();
  TORCH_CHECK(size > 0, "median() input tensor cannot be empty");

  // Sort input tensor to efficiently query for median element
  Tensor sorted = std::get<0>(self.flatten().sort());

  if (!ignore_nan) {
    // For torch.median return either the middle element or nan (sorted as
    // largest) if there are any
    int64_t k = (size - 1) / 2;
    return at::where(sorted[-1].isnan(), sorted[-1], sorted[k]);
  } else {
    // For torch.nanmedian return the middle element among the non-nan values
    Tensor k = ((size - 1) - sorted.isnan().sum()) / 2;
    return sorted[k.toType(kLong)];

On the other hand, CuPyv8.5’s median calls partition internally:

github.com

cupy/cupy/blob/6d91f956ee3c83d07fe0548268345d46fa319ccd/cupy/core/_routines_statistics.pyx#L410-L414


if axis is None:
    part = part.ravel()
    part.partition(kth)
else:
    part.partition(kth, axis=axis)

In NumPy, sort's order is O(n^2) in the worst case while partition, O(n) according to https://stackoverflow.com/a/43589598/8335699. And if this is true to PyTorch and CuPy, the gap you found can be possible.