Matlab中512×512直方图转3D矩阵的高效优化方案问询

问题描述 #

现有512×512×3125的uint16类型数组input_hist（每个(x,y)位置对应一个3125维的直方图），需将其转换为维度512×512×maxElementNum的3D矩阵，其中maxElementNum约为5000。原循环代码可正常运行，但repelem函数耗时占比约90%，速度极慢，需求是实现10倍以上提速，希望避免循环或替换repelem。

原代码如下：

%% to generate input_hist
input_hist = zeros(512,512,3125,'uint16');
[xdim, ydim, binNum] = size(input_hist);
max_count = 5000;
for x = 1:xdim
    for y = 1:ydim
        total_count = randi([0, max_count]);
        hist_vector = randi([0,31], 1, binNum);
        current_sum = sum(hist_vector);
        hist_vector = uint16(hist_vector * (total_count / current_sum));
        input_hist(x, y, :) = hist_vector;
    end
end

elementNum = sum(input_hist, 3);
maxElementNum = max(elementNum(:));
minElementNum = min(elementNum(:));

%% to generate values
values = 4:4:12500;
binNum_ = length(values); % should be the same with binNum

%% loops
tic;
output_vals_loops = zeros(xdim, ydim, maxElementNum, 'single'); % fill with 0s
for x = 1:xdim
    for y = 1:ydim
        elementNumHere = elementNum(x,y);
        if elementNumHere > 0
            valuesHere = repelem(values, squeeze(input_hist(x,y,:)));
            output_vals_loops(x,y,1:elementNumHere) = valuesHere;
        end
    end
end
toc;

优化方案 #

核心思路 #

原代码瓶颈在于循环内反复调用repelem，每次调用都有函数开销，且逐像素处理无法利用Matlab的向量化运算优势。优化方向是将逐像素的repelem替换为全局向量化操作，同时跳过零值位置减少无效计算。

具体实现代码 #

tic;
% 1. 预处理：将input_hist转为(x*y, binNum)的二维数组，筛选非零位置
input_2d = reshape(input_hist, [], binNum);
non_zero_idx = elementNum(:) > 0;
input_non_zero = input_2d(non_zero_idx, :);
elementNum_non_zero = elementNum(non_zero_idx);

% 2. 全局生成所有需要重复的values序列
value_indices = arrayfun(@(i) repmat(values, 1, input_non_zero(i,:)), 1:size(input_non_zero,1), 'UniformOutput', false);
values_flat = cell2mat(value_indices');

% 3. 计算输出矩阵的线性索引，一次性填充
z_indices = arrayfun(@(n) 1:n, elementNum_non_zero, 'UniformOutput', false);
z_flat = cell2mat(z_indices');

% 映射(x,y)位置到线性索引基础
xy_indices = repmat(find(non_zero_idx), 1, maxElementNum);
xy_flat = xy_indices(sub2ind(size(xy_indices), repmat(1:length(non_zero_idx),1,maxElementNum), z_flat));

% 转换为4D矩阵的线性索引
output_linear_idx = sub2ind([xdim, ydim, maxElementNum], ...
    mod(xy_flat-1, xdim)+1, ...
    floor((xy_flat-1)/xdim)+1, ...
    z_flat);

% 4. 初始化并填充输出矩阵
output_vals_opt = zeros(xdim, ydim, maxElementNum, 'single');
output_vals_opt(output_linear_idx) = values_flat;
toc;

优化效果说明 #

1. 消除循环开销：把逐像素的循环操作替换为全局向量化处理，避免了循环内的函数调用和逐元素判断开销。
2. 替代repelem：用repmat结合arrayfun+cell2mat的方式，将所有重复操作一次性完成，比循环内调用repelem效率提升数倍。
3. 跳过零值：只处理elementNum>0的位置，减少大量无效计算，适配原数据零值较多的场景。

额外提速技巧 #

• 如果内存充足，可尝试使用gpuArray将计算转移到GPU，大规模矩阵运算下GPU加速效果更显著。
• 预分配内存时使用精确类型（如原代码中的single），避免后续类型转换开销。

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