如何将dlib::array2d<dlib::rgb_pixel>转为torch at::Tensor并解决数值异常？

问题根源 #

你遇到的超大值（1e+38、1e+40）和模型输出NaN，核心不是slice操作的问题，而是数据类型不匹配、内存布局错误、内存拷贝逻辑漏洞导致的：

• dlib的rgb_pixel是unsigned char类型（取值0-255），如果直接拷贝到float张量却没做数值缩放，或者拷贝时指针越界读取了未初始化内存，就会出现异常大值
• PyTorch模型通常期望NCHW格式的输入（批量、通道、高、宽），但dlib的array2d是HWC格式（高、宽、通道），通道顺序默认是RGB，若和模型训练时的输入格式（比如BGR）不一致，也会引发后续计算异常
• 直接用内存指针强制转换而忽略张量的步长（stride），会导致内存读取错位，产生随机垃圾值

正确的转换实现 #

下面是经过验证的转换函数，包含格式转换、数值归一化、通道调整的完整逻辑：

#include <torch/torch.h>
#include <dlib/image_processing.h>
#include <dlib/gui_widgets.h>
#include <dlib/image_io.h>

torch::Tensor dlib_image_to_tensor(const dlib::array2d<dlib::rgb_pixel>& img) {
    // 获取图像的高、宽
    int height = img.nr();
    int width = img.nc();
    int channels = 3;

    // 1. 创建HWC格式的临时张量，数据类型匹配rgb_pixel的unsigned char
    auto tensor_hwc = torch::empty({height, width, channels}, torch::kUInt8);
    auto tensor_data = tensor_hwc.data_ptr<unsigned char>();

    // 2. 逐像素拷贝dlib图像数据到张量
    for (int h = 0; h < height; ++h) {
        for (int w = 0; w < width; ++w) {
            const dlib::rgb_pixel& pixel = img(h, w);
            // RGB通道顺序，对应张量的第三个维度
            tensor_data[h * width * channels + w * channels + 0] = pixel.red;
            tensor_data[h * width * channels + w * channels + 1] = pixel.green;
            tensor_data[h * width * channels + w * channels + 2] = pixel.blue;
        }
    }

    // 3. 转换为模型需要的格式：
    //    - 转成float32类型并归一化（和训练时的预处理一致，这里以除以255为例）
    //    - 从HWC转成CHW，再添加batch维度变成NCHW
    auto tensor = tensor_hwc.to(torch::kFloat32) / 255.0f;
    tensor = tensor.permute({2, 0, 1}).unsqueeze(0);

    // 如果模型训练时用的是BGR格式，需要交换通道顺序
    // tensor = tensor.index_select(0, torch::tensor({2, 1, 0}, torch::kLong));

    return tensor;
}

关键细节说明 #

• 数据类型对齐：先用torch::kUInt8类型的张量承接dlib的像素数据，避免直接转float时出现类型转换错误
• 归一化操作：必须将0-255的像素值缩放到模型训练时的输入范围（比如0-1或[-1,1]），否则超大的输入值会导致模型激活函数溢出，最终输出NaN
• 格式转换：通过permute将HWC转为CHW，再用unsqueeze(0)添加batch维度，匹配PyTorch模型的输入要求
• 通道顺序检查：如果你的模型是用OpenCV等工具训练（默认BGR），需要打开代码中注释的通道交换逻辑

验证方法 #

转换后可以打印张量的基本信息，确认数值范围正常：

auto tensor = dlib_image_to_tensor(img);
std::cout << "Tensor shape: " << tensor.sizes() << std::endl;
std::cout << "Min value: " << tensor.min().item<float>() << std::endl;
std::cout << "Max value: " << tensor.max().item<float>() << std::endl;

正常情况下，min应该接近0，max接近1（如果用除以255的归一化），没有超大值。此时再输入模型，前向传播结果就不会出现NaN了。

内容的提问来源于stack exchange，提问作者Ulrich Von Rekkenin