时间戳

首页 > 解决方案 > Julia 没有使用所有可用的 CPU 线程

multithreading - Julia 没有使用所有可用的 CPU 线程

问题描述

我在 Julia 中进行了一些计算，发现它的运行速度明显比 numpy慢（大约25倍！）。

然后我意识到 Julia 在我的 PC 上总共96 个CPU 线程（48 个物理内核）中只使用了8个线程，而 numpy 使用70多个线程似乎没有问题。

即使返回，使用$julia --thread 96 参数运行 Julia也没有任何区别。julia> Threads.nthreads()96

此外，结果有点令人失望的是，我怀疑 Julia 使用所有 96 个线程仍然可能无法与 numpy 的速度相匹配。

这是 Julia 代码。我只是测量julia> @time calc_fid(mat_a, mat_b)给我平均90秒的时间。

using Statistics
using LinearAlgebra

function calc(A::Array{Float32,2}, B::Array{Float32,2})
    μ_A = mean(A, dims=2)
    μ_B = mean(B, dims=2)
    σ_A = cov(A, dims=2)
    σ_B = cov(B, dims=2)

    ssdiff = sum((μ_A - μ_B).^2)
    covmean = sqrt(σ_A * σ_B)
    res = ssdiff + tr(σ_A .+ σ_B .- 2.0 * covmean)
    
    return res
end

这是平均需要大约3.5秒的 numpy 代码。用time.perf_counter()

import numpy as np

from numpy import cov
from numpy import trace
from scipy.linalg import sqrtm

def calc(A, B):
    mu_A = A.mean(axis=0)
    mu_B = B.mean(axis=0)
    sigma_A = cov(A, rowvar=False)
    sigma_B = cov(B, rowvar=False)

    ssdiff = np.sum((mu_A - mu_B) ** 2.0)
    covmean = sqrtm(sigma_A.dot(sigma_B))
    res = ssdiff + trace(sigma_A + sigma_B - 2.0 * covmean)

    return res

任何建议/解释将不胜感激！

标签： multithreadingperformancenumpyjuliabenchmarking

解决方案

您的 Julia 示例代码实际上并未使用 Julia 的多线程功能。

在 Julia 中，如果你想要多线程，通常需要显式启用它。仅仅凭借多线程启动 Julia，类似mean或不会自动成为多线程的函数；cov如果您希望这些函数多线程，您将需要编写它们的多线程版本，或者使用已经这样做的包。

在您的 Julia 示例代码中，目前编写的唯一将多线程的东西是线性代数，因为这会回到BLAS（确实可能是 numpy），它有自己的多线程系统，与 Julia 的多线程完全分离（BLAS 是用 Fortran 编写）。不过，这也可能没有使用所有 96 个线程，并且在您的情况下几乎可以肯定默认为 8。可以检查 Julia 中 BLAS 使用的线程数

using LinearAlgebra
BLAS.get_num_threads()

或类似地设置

BLAS.set_num_threads(n)

推荐阅读