关于Julia的GPU内存管理

综合讨论区
1

大佬们，使用GPU循环计算，循环次数过多之后就会出现内存不足死掉的情况。将所有涉及GPU的变量在循环末尾设置 = nothing 也不行，官方手册对GPU内存管理方面介绍不多，请问有没有其他内存管理方法。

1 个赞
2

奉上代码，请各位指点

for col_start in 1:100:1500
    col_end = min(col_start+99, 1500)
    col_range = col_start:col_end
    
     # 创建 SeisChannel 对象
    R = []
    for col in col_range
        S = SeisChannel(
            id = " ",
            name = " ",
            fs = 500.0,
            x = dasdata[col, :],
            t = t
        )
        Rsin = RawData(S, cc_len, cc_step) |> gpu
        push!(R, Rsin)
    end
    
    # 预处理
    detrend!.(R)
    taper!.(R)
    bandpass!.(R, 1, 30, zerophase=true)
    onebit!.(R)
    FFT = rfft.(R)
    whiten!.(FFT, 1, 30)
    
    # 初始化存储互相关结果的字典
    Cm = correlate(FFT[1], FFT[2], 4) |> cpu
    cross_corr_results = Dict{Tuple{Int, Int}, typeof(Cm)}()
    
    # 计算互相关并存储结果
    for i in 1:arraylt-1
        for j in i+1:arraylt
            C = correlate(FFT[i], FFT[j], 4)
            abs_max!(C)
            cross_corr_results[(i, j)] = C |> cpu
        end
    end
    
    # 初始化总叠加结果矩阵
    Cs_total = zeros(4001, arraylt-1)
    
    # 叠加互相关结果
    Threads.@threads for i in 1:arraylt-1
        for j in i+1:arraylt
            C = cross_corr_results[(i, j)]
            pws!(C)
            Cs_total[:, j-i] += C.corr
        end
    end

    channel_start = 2900 + 1500 + col_start - 1
    channel_end = 2900 + 1500 + col_end - 1
    # 保存叠加互相关结果
    save_folder = "  "
    if !isdir(save_folder)
        mkpath(save_folder)
    end
    save_path = "  "
    writedlm(save_path, Cs_total, ',')

    R = nothing
    FFT = nothing
    cross_corr_results = nothing
    CUDA.synchronize()
    GC.gc()

end
3

补充信息：这个代码单GPU运行是没有问题的，问题可能出在两块GPU的并行上，我的并行方法是将原始数据拆分为两半，通过Threads.@spawn运行两个函数，每个函数中调用CUDA.device!()指定使用GPU。