对数组操作 并行化

有一个三维的tensor,我想对每个tensor[i, j, :]序列都进行去重复求其均值,写成函数如下:

function eliminate(tensor::AbstractArray{T, 3}) where T
    dims_1, dims_2, dims_3 = size(tensor)
    mean_array = similar(tensor, dims_1, dims_2)

    @inbounds @views for i in 1: dims_1, j in 1: dims_2
        mean_array[i, j] = unique(tensor[i, j, :]) |> mean
    end

    return mean_array
end

我想对这段代码进行并行化,提高运行速度,第一次接触并行化, 请问有什么好的建议吗?

我后来改写了一下:

function eliminate2(tensor::AbstractArray{T, 3}) where T
    dims_1, dims_2, dims_3 = size(tensor)
    mean_array = zeros(T, dims_1, dims_2)

    @inbounds for i in 1: dims_1, j in 1: dims_2
        hash_tab = Set{T}()
        clock = 0
        for k in 1: dims_3
            if tensor[i, j, k] ∉ hash_tab
                push!(hash_tab, tensor[i, j, k])
                mean_array[i, j] += tensor[i, j, k]
            clock += 1
            end
        end
        mean_array[i, j] /= clock
    end

    return mean_array
end

这个版本的函数比上面的版本快了不少,但还是不够。我想说如果对这段代码进行并行化处理的话,是不是可以加速。但是由于涉及到对数组的操作,一头雾水。


-    @inbounds for i in 1: dims_1, j in 1: dims_2
+   Threads.@threads for ind in CartesianIndices((1: dims_1, 1: dims_2))
+        i, j = ind.I

这里你对每个 element 的操作都是独立的,直接上 threads就可以了

1赞

刚想到了,如果你用 GPU 加速的话,应该也是可以的,先 sort 下,然后写个kernel的函数,遇到重复的就跳过,应该可以轻松实现加速

1赞

好的,非常感谢,我去尝试一下

京ICP备17009874号-2