Median of multiple arrays in Julia

user_4744283 0
Viewed 428

It's related to this question

I want to know how to calculate median along specific dimension on huge array, for example with size (20, 1920, 1080, 3). I not sure whether there is any practical purpose but I just wanted to check how well median works in Julia.

It takes ~0.5 seconds to calculate medians on (3,1920,1080,3) with numpy. It works very fast on zeros array (less than 2 seconds on (120, 1920, 1080,3)) and works not so fast but fine on real images (20 seconds on (120, 1920, 1080,3)).

Python code:

import cv2
import sys
import numpy as np
import time

ZEROES=True
N_IMGS=20

print("n_imgs:", N_IMGS)
print("use dummy data:", ZEROES)

imgs_paths = sys.argv[1:]
imgs_paths.sort()
imgs_paths_sparse = imgs_paths[::30]

imgs_paths = imgs_paths_sparse[N_IMGS]

if ZEROES:
    imgs_arr = np.zeros((N_IMGS,1080,1920,3), dtype=np.float32)
else:
    imgs = map(cv2.imread, imgs_paths)
    imgs_arr = np.array(list(imgs), dtype=np.float32)

start = time.time()
imgs_median = np.median(imgs_arr, 0)
end = time.time()
print("time:", end - start)
cv2.imwrite('/tmp/median.png', imgs_median)

In julia I can only calculate median of (3, 1920, 1080,3). After that my earlyoom process kills julia process because of huge amount of used memory.

I tried approach similar to what I tried first on max:

function median1(imgs_arr)
    a = imgs_arr
    b = reshape(cat(a..., dims=1), tuple(length(a), size(a[1])...))
    imgs_max = Statistics.median(b, dims=1)
    return imgs_max
end

Or even more simple case:

import Statistics
a = zeros(3,1080,1920,3)
@time Statistics.median(a, dims=1)
 10.609627 seconds (102.64 M allocations: 2.511 GiB, 3.37% gc time)
...

So, it takes 10 seconds vs 0.5 seconds on numpy. I have only 4 CPU cores and it's not simply parallelization.

Is there more or less simple way to optimize it somehow?

Or at least take slices and compute it one-by-one without overuse of memory?

julia median arrays
2 Answers

Try JuliennedArrays.jl

julia> a = zeros(3,1080,1920,3);

julia> using JuliennedArrays

julia> @time map(median, Slices(a,1));
  0.822429 seconds (6.22 M allocations: 711.915 MiB, 20.15% gc time)

As Stefan commented below, the built in median does the same thing, but much slower

julia> @time median(a, dims=1);
  7.450394 seconds (99.80 M allocations: 2.368 GiB, 4.47% gc time)

at least as of julia> VERSION v"1.5.0-DEV.876"

user_19147910

It's hard to know if the fact that the images are loaded separately is a key part of the problem here or not since the setup for the problem in Julia is missing and it's a bit hard for Julia programmers to follow the Python setup or know how much we need to match it. You either need to:

  1. Load or move the image data so that they are, in fact, part of the same array and then take the median of that;

  2. Make a set of spatially unrelated values in different arrays abstractly behave as though they are part of a single array and then take the median of that collection via a method that's generic enough to handle this abstraction.

Fredrik's answer implicitly assumes that you have already loaded the image data so that they're all part of the same contiguous array. If that's the case, however, then you don't even need JuliennedArrays, you can just use the median function from the Statistics stdlib:

julia> a = rand(3, 1080, 1920, 3);

julia> using Statistics

julia> median(a, dims=1)
1×1080×1920×3 Array{Float64,4}:
[:, :, 1, 1] =
 0.63432  0.205958  0.216221  0.571541  …  0.238637  0.285947  0.901014

[:, :, 2, 1] =
 0.821851  0.486859  0.622313  …  0.917329  0.417657  0.724073

If you can load the data like this, it's the best approach—this is by far the most efficient representation of a bunch of same-sized images and makes vectorize operations across images easy and efficient. The first dimension is the most efficient one to do operations across because Julia is column-major, so the first dimension (columns) is stored contiguously.

The best way to get the images into contiguous memory is to pre-allocate an uninitialized array of the right type and dimensions and then read the data into the array using some in-place API. For some reason your Julia code appears to have loaded the images as a vector of individual arrays while your Python code seems to have loaded all of the images into a single array?

The approach of reshaping and concatenating is an extreme case of the second approach where you move all of the data all at once before then applying a vectorized median operation. Obviously, that involves moving a lot of data around, which is pretty inefficient.

Due to memory locality, it may be more efficient to copy a single slice of the data into a temporary array and compute the median of that. That can be done pretty easily with an array comprehension:

julia> v_of_a = [rand(1080, 1920, 3) for _ = 1:3]
3-element Array{Array{Float64,3},1}:
 [0.7206652600431633 0.7675119703509619 … 0.7117084561740263 0.8736518021960584; 0.8038479801395197 0.3159392943734012 … 0.976319025405266 0.3278606124069767; … ; 0.7424260315304789 0.4748658164109498 … 0.9942311708400311 0.37048961459068086; 0.7832577306186075 0.13184454935145773 … 0.5895094390350453 0.5470111170897787]

[0.26401298651503025 0.9113932653115289 … 0.5828647778524962 0.752444909740893; 0.5673144007678044 0.8154276504227804 … 0.2667436824684424 0.4895443896447764; … ; 0.2641913584303701 0.16639100493266934 … 0.1860616855126005 0.04922131616483538; 0.4968214514330498 0.994935452055218 … 0.28097239922248685 0.4980189891952156]

julia> [median(a[i,j,k] for a in v_of_a) for i=1:1080, j=1:1920, k=1:3]
1080×1920×3 Array{Float64,3}:
[:, :, 1] =
 0.446895  0.643648  0.694714   …  0.221553   0.711708   0.225268
 0.659251  0.457686  0.672072      0.731218   0.449915   0.129987
 0.573196  0.328747  0.668702      0.355231   0.656686   0.303168
 0.243656  0.702642  0.45708       0.23415    0.400252   0.482792
user_6592480
Related