xarray слишком медленный для критически важного для производительности кода

Я планировал широко использовать xarray в некотором научном коде, который я пишу. Пока это делает код очень элегантным, но я думаю, что мне придется отказаться от него, так как цена производительности слишком высока.

Вот пример, который создает два массива и перемножает их части вместе, используя xarray (с несколькими схемами индексации) и numpy. Я использовал num_comp=2 и num_x=10000:

Line #      Hits     Time   Per Hit   % Time  Line Contents
 4                                           @profile
 5                                           def xr_timing(num_comp, num_x):
 6         1         4112   4112.0     10.1      da1 = xr.DataArray(np.random.random([num_comp, num_x]).astype(np.float32), dims=['component', 'x'], coords={'component': ['a', 'b'], 'x': np.linspace(0, 1, num_x)})
 7         1          438    438.0      1.1      da2 = da1.copy()
 8         1         1398   1398.0      3.4      da2[:] = np.random.random([num_comp, num_x]).astype(np.float32)
 9         1         7148   7148.0     17.6      da3 = da1.isel(component=0).drop('component') * da2.isel(component=0).drop('component')
10         1         6298   6298.0     15.5      da4 = da1[dict(component=0)].drop('component') * da2[dict(component=0)].drop('component')
11         1         7541   7541.0     18.6      da5 = da1.sel(component='a').drop('component') * da2.sel(component='a').drop('component')
12         1         7184   7184.0     17.7      da6 = da1.loc[dict(component='a')].drop('component') * da2.loc[dict(component='a')].drop('component')
13         1         6479   6479.0     16.0      da7 = da1[0, :].drop('component') * da2[0, :].drop('component')

15                                           @profile
16                                           def np_timing(num_comp, num_x):
17         1         1027   1027.0     50.2      da1 = np.random.random([num_comp, num_x]).astype(np.float32)
18         1          977    977.0     47.8      da2 = np.random.random([num_comp, num_x]).astype(np.float32)
19         1           41     41.0      2.0      da3 = da1[0, :] * da2[0, :]

Самое быстрое умножение xarray занимает примерно в 150 раз больше времени, чем версия numpy. Это всего лишь одна из операций в моем коде, но я обнаружил, что большинство из них во много раз медленнее, чем эквивалент numpy, что досадно, поскольку xarray делает код намного понятнее. Я делаю что-то неправильно?

Обновление: даже da1[0, :].values ​​* da2[0, :].values ​​(при котором теряются многие преимущества использования xarray) занимает 2464 единицы времени.

Я использую xarray 0.9.6, pandas 0.21.0, numpy 1.13.3 и Python 3.5.2.

Обновление 2: по просьбе @Maximilian, вот повторный запуск с num_x=1000000:

Line #      Hits   Time    Per Hit   % Time  Line Contents
# xarray
 9         5       408596  81719.2     11.3      da3 = da1.isel(component=0).drop('component') * da2.isel(component=0).drop('component')
10         5       407003  81400.6     11.3      da4 = da1[dict(component=0)].drop('component') * da2[dict(component=0)].drop('component')
11         5       411248  82249.6     11.4      da5 = da1.sel(component='a').drop('component') * da2.sel(component='a').drop('component')
12         5       411730  82346.0     11.4      da6 = da1.loc[dict(component='a')].drop('component') * da2.loc[dict(component='a')].drop('component')
13         5       406757  81351.4     11.3      da7 = da1[0, :].drop('component') * da2[0, :].drop('component')
14         5        48800   9760.0      1.4      da8 = da1[0, :].values * da2[0, :].values

# numpy
20         5        37476   7495.2      2.9      da3 = da1[0, :] * da2[0, :]

Как и ожидалось, разница в производительности существенно уменьшилась (теперь всего в 10 раз медленнее), но я все равно рад, что проблема будет упомянута в следующем выпуске документации, поскольку даже такая разница может удивить некоторых людей.


person user3708067    schedule 08.11.2017    source источник
comment
Можете повторить сравнение на больших массивах? Разница в производительности может значительно уменьшиться, так как сама вычислительная часть должна иметь одинаковую скорость.   -  person Maximilian    schedule 09.11.2017


Ответы (1)


Да, это известное ограничение для xarray. Код, чувствительный к производительности, который использует небольшие массивы, намного медленнее для xarray, чем NumPy. Я написал об этом новый раздел в нашей документации для следующей версии: http://xarray.pydata.org/en/stable/computation.html#wrapping-custom-computation

В основном у вас есть два варианта:

  1. Напишите свой код, чувствительный к производительности, в развернутых массивах, а затем оберните их обратно в структуры данных xarray. Xarray v0.10 имеет новую вспомогательную функцию (apply_ufunc), которая немного упрощает эту задачу. См. ссылку выше, если вы заинтересованы в этом.
  2. Используйте что-то кроме xarray/Python для выполнения ваших вычислений. Это также может иметь смысл, потому что сам Python добавляет значительные накладные расходы. AxisArrays.jl Джулии выглядит интересно, хотя я сам не пробовал.

Я предполагаю, что вариант 3 будет состоять в том, чтобы переписать сам xarray на C++ (например, поверх xtensor), но это было бы гораздо больше вовлечено!

person shoyer    schedule 08.11.2017