Как я могу оптимизировать этот огромный цикл, использующий libtorch и numcpp ⇐ C++

[Anonymous]

Я пишу код для суперразрешения с помощью cpp, и в этом примере я использую libtorch и numcpp.
Мне нужно вычислить информацию о пикселях img после масштабирования (масштаб равен 2 в моем примере), который будет тратить время цикла new_img_width * new_img_height. Однако я также должен получить новую информацию из исходного img (то есть я добавляю внутренний цикл с old_img_width * old_img_height раз во внешнем цикле) и общее время циклов будет огромным (new_img_width * new_img_height * old_img_width * old_img_height), мне нужно провести некоторую оптимизацию циклов (внутренний цикл тратит всего 0,005 с), но все это время все еще долго, потому что внешние циклы огромны, как я могу оптимизировать свои код?
Дополнительная информация к коду:
1.merging_t — это класс, содержащий только некоторые скалярные данные
2.nc::NdArray — это numpy.ndarray для cpp
3.другие функции в кодах просто выполняют быстрые математические вычисления

Код: Выделить всё

    void merge(nc::NdArray& comp_img, at::Tensor& alignments, at::Tensor& covs, nc::NdArray r, at::Tensor& num, at::Tensor& den, int options, params::merging_t& params)
{
float scale = params.scale;
nc::NdArray CFA_pattern = params.exif.CFA_Pattern;
bool bayer_mode = params.mode == "bayer";
bool iso_kernel = params.kernel == "iso";
int tile_size = params.tuning.tileSize;
int input_size_y = comp_img.shape().rows, input_size_x =comp_img.shape().cols;
int64_t align_y = alignments.sizes()[0], align_x = alignments.sizes()[1], align_c = alignments.sizes()[2];
int64_t num_y = num.sizes()[0], num_x = num.sizes()[1], num_c = num.sizes()[2];
int64_t den_y = den.sizes()[0], den_x = den.sizes()[1], den_c = den.sizes()[2];
int64_t covs_y = covs.sizes()[0], covs_x = covs.sizes()[1];
float* num_data = new float[num_y*num_x*num_c]();
float* den_data = new float[den_y * den_x * den_c]();
float* align_data = new float[align_y * align_x * align_c];
float* covs_data = new float[covs_y * covs_x * 2 * 2];
std::memcpy(covs_data, covs.data_ptr(), covs.numel() * sizeof(float));
std::memcpy(align_data, alignments.data_ptr(), sizeof(float) * alignments.numel());
//H W C
for (int output_pixel_idy = 0; output_pixel_idy < num_y; ++output_pixel_idy)
{
for (int output_pixel_idx = 0; output_pixel_idx < num_x; ++output_pixel_idx)
{

nc::NdArray acc;
nc::NdArray val;
int n_channels;
if (bayer_mode)
{
n_channels = 3;
acc = nc::zeros(1, 3);
val = nc::zeros(1, 3);
}
else
{
n_channels = 1;
acc = nc::zeros(1, 1);
val = nc::zeros(1, 1);
}
nc::NdArray local_CFA = CFA_pattern;
float coarse_ref_sub_pos[2] = { output_pixel_idy / scale,output_pixel_idx / scale };

nc::NdArray local_optical_flow = nc::zeros(1, 2);
int patch_idy = coarse_ref_sub_pos[0] / tile_size;
int patch_idx = coarse_ref_sub_pos[1] / tile_size;

local_optical_flow[0] = align_data[patch_idy * align_x * align_c + patch_idx * align_c + 0];
local_optical_flow[1] = align_data[patch_idy * align_x * align_c + patch_idx * align_c + 1];

nc::NdArray patch_center_pos = { coarse_ref_sub_pos[0] + local_optical_flow[1] ,coarse_ref_sub_pos[1] + local_optical_flow[0] };
int x_r=0, y_r=0;
y_r = utils::clamp(std::round(coarse_ref_sub_pos[0]), 0, r.shape().rows - 1);
x_r = utils::clamp(std::round(coarse_ref_sub_pos[1]), 0, r.shape().cols - 1);
float local_r = r[y_r, x_r];

if (!(patch_center_pos[1] >= 0 && patch_center_pos[1] < input_size_x && patch_center_pos[0] >= 0 && patch_center_pos[0] < input_size_y))
{
output_pixel_idx = num_x;
output_pixel_idy = num_y;
break;
}

nc::NdArray interpolated_cov(nc::Shape(2, 2)), cov_i(nc::Shape(2, 2));
float close_covs[2][2][2][2] = {};
nc::NdArray  grey_pos(nc::Shape(1, 2));

if (!iso_kernel)
{
if (bayer_mode)
{
grey_pos[0] = (coarse_ref_sub_pos[0] - 0.5) / 2;
grey_pos[1] = (coarse_ref_sub_pos[1] - 0.5) / 2;
}
else
{
grey_pos[0] = coarse_ref_sub_pos[0];
grey_pos[1] = coarse_ref_sub_pos[1];
}
int floor_x = max(std::floor(grey_pos[1]), 0);
int floor_y = max(std::floor(grey_pos[0]), 0);
int ceil_x = min(floor_x + 1, covs.sizes()[1] - 1);
int ceil_y = min(floor_y + 1, covs.sizes()[0] - 1);
for (int i = 0; i < 2; ++i)
{
for (int j = 0; j < 2; ++j)
{
close_covs[0][0][i][j] = covs_data[floor_y * covs_x * 2 * 2 + floor_x * 2 * 2 + i * 2 + j];
close_covs[0][1][i][j] = covs_data[floor_y * covs_x * 2 * 2 + ceil_x * 2 * 2 + i * 2 + j];
close_covs[1][0][i][j] = covs_data[ceil_y * covs_x * 2 * 2 + floor_x * 2 * 2 + i * 2 + j];
close_covs[1][1][i][j] = covs_data[ceil_y * covs_x * 2 * 2 + ceil_x * 2 * 2 + i * 2 + j];
}
}
interpolated_cov = linalg::interpolate_cov(close_covs, grey_pos);
cov_i = linalg::invert_2x2(interpolated_cov);
}

int center_x = std::round(patch_center_pos[1]);
int center_y = std::round(patch_center_pos[0]);
for (int i = -1; i < 2; ++i)
{
for (int j = -1; j < 2; ++j)
{
for (int pixel_idx = center_x + j; pixel_idx < input_size_x; ++pixel_idx)
{
for (int pixel_idy = center_y + i; pixel_idy < input_size_y; ++pixel_idy)
{
int channel = bayer_mode ? local_CFA(pixel_idy % 2, pixel_idx % 2) : 0;
float c =comp_img(pixel_idy, pixel_idx);
float dist_x = pixel_idx - patch_center_pos[1];
float dist_y = pixel_idy - patch_center_pos[0];
float y;
if (iso_kernel)
{
y = max(0, 2 * (dist_x * dist_x + dist_y * dist_y));
}
else
{
y = max(0, linalg::quad_mat_prod(cov_i, dist_x, dist_y));
}

float w = std::exp(-0.5 * y);
val[channel] += c * w * local_r;
acc[channel] += w * local_r;
}
}
}
}
for (int chan = 0; chan < n_channels; ++chan)
{
num_data[output_pixel_idy * num_x * num_c + output_pixel_idx * num_c + chan] += val[chan];
den_data[output_pixel_idy * num_x * num_c + output_pixel_idx * num_c + chan] += acc[chan];
}
}
}
num = torch::from_blob(num_data, { num_y,num_x,num_c }, c10::kFloat).clone();
den = torch::from_blob(den_data, { den_y,den_x,den_c }, c10::kFloat).clone();
delete[] num_data;
delete[] den_data;
delete[] covs_data;
delete[] align_data;
}
}

в моем коде первой редакции я использую оператор перегрузки индекса torch::Tensor (т.е. оператор[]), изменяя данные Tensor напрямую, теперь я знаю, что это очень медленно, поэтому я использую исходный указатель для изменения данных и построения из них тензора во втором издании (это пример), правда в том, что общее время все еще очень велико, хотя я заменяю все места, где могу использовать исходный указатель.
Я пробовал omp.h, особого эффекта не дало.
может быть, cuda(GPU) — хорошее направление, но я AMD

Подробнее здесь: https://stackoverflow.com/questions/781 ... and-numcpp