如何用Julia高效读取4000张TinyMNIST图片并转为Matrix{Float64}

我需要用Julia读取TinyMNIST数据集的10个目录下共4000张图片，并将其数值整理为Matrix{Float64}格式。目前编写的代码如下，基准测试显示耗时较长，尝试用TensorCast.jl优化reduce(hcat, c)效果不明显，求加快处理速度的方法。

用户原代码 #

using FileIO, Images, BenchmarkTools

const paths::Vector{String} = readdir(
    raw"<PATH>\TinyMNIST\train"; join=true
)

function ret4000(paths::Vector{String})::Matrix{Float64}
    reduce(hcat, _ret400.(paths))
end

function _ret400(path::String)::Matrix{Float64}
    # `set1` is a Vector of 400 images' path
    set1::Vector{String} = readdir(path; join=true);
    # `b` is a Vector of 400 matrixes of size 28x28
    b::Vector{Matrix{Float64}} = [Float64.(x) for x=FileIO.load.(set1)]
    # `c` is a Vector that contains 400 matrix of size 784x1
    c::Vector{Matrix{Float64}} = [reshape(x, :, 1) for x=b]
    # horizontal concatenation of all 400 matrixes of size 784x1 to form a matrix of size 784x400
    reduce(hcat, c)
end

基准测试结果 #

julia> @benchmark ret4000($paths)
BenchmarkTools.Trial: 1 sample with 1 evaluation.
 Single result which took 5.187 s (0.25% GC) to evaluate,
 with a memory estimate of 148.41 MiB, over 1222457 allocations.

julia> typeof(ret4000(paths))
Matrix{Float64} (alias for Array{Float64, 2})

julia> size(ret4000(paths))
(784, 4000)

优化方案 #

1. 预分配内存，消除临时数组开销 #

原代码多次创建中间临时数组（b、c及hcat的中间结果），带来大量内存分配和GC负担。直接预分配最终大小的矩阵，逐个填充数据：

function ret4000_opt(paths::Vector{String})::Matrix{Float64}
    # 预分配最终矩阵：784行，4000列
    result = Matrix{Float64}(undef, 784, 4000)
    col_idx = 1
    for path in paths
        set1 = readdir(path; join=true)
        for img_path in set1
            # 加载图片后直接转一维数组填充
            img = Float64.(FileIO.load(img_path))
            result[:, col_idx] .= vec(img)
            col_idx += 1
        end
    end
    return result
end

该方式复用预分配内存，大幅减少分配次数与GC时间。

2. 多线程并行加载图片 #

图片加载属于IO密集型操作，利用Julia多线程并行处理（启动Julia时需加-t auto开启多线程）：

function ret4000_parallel(paths::Vector{String})::Matrix{Float64}
    result = Matrix{Float64}(undef, 784, 4000)
    # 提前收集所有图片路径，避免重复目录IO
    all_img_paths = vcat([readdir(p; join=true) for p in paths]...)
    # 并行处理每张图片
    Threads.@threads for i in 1:length(all_img_paths)
        img = Float64.(FileIO.load(all_img_paths[i]))
        result[:, i] .= vec(img)
    end
    return result
end

并行处理能充分利用多核CPU，减少IO等待时间，显著提升整体速度。

3. 优化图片类型转换 #

FileIO.load返回的图片包含元数据，可通过rawview和channelview直接提取原始像素值，减少冗余转换：

# 高效提取灰度图像素值示例
img = load(img_path)
img_data = Float64.(rawview(channelview(img)))

4. 消除冗余reshape操作 #

原代码中reshape(x, :, 1)再hcat的步骤完全可以用vec(x)替代，直接将一维数组填充到预分配矩阵中，省去多次拼接的开销。

内容的提问来源于stack exchange，提问作者Shayan