C++结合OpenCV调用traced模型(.pt)输出与Python端结果存在差异如何解决

我基于EnlightenGAN重新训练了模型，为实现C端部署，对模型做trace处理得到.pt文件，计划在搭载libTorch v1.6的C应用中运行。但实际推理后，C++端输出结果与Python端执行同一traced模型的结果存在轻微差异。
该模型需要输入RGB张量与注意力图张量两个参数，注意力图的作用是告知模型需要执行对比度增强的图像区域。

Python端推理代码 #

def getTransform():
    transform_list = []
    transform_list += [transforms.ToTensor(),
                       transforms.Normalize((0.5, 0.5, 0.5),
                                            (0.5, 0.5, 0.5))]
    return transforms.Compose(transform_list)

def convertToCV(tensor):
    
    tensor = torch.squeeze(tensor)
    tensor = tensor.cpu().float().detach()
    tensor = torch.unsqueeze(tensor, 0)
    tensor = tensor.permute(1, 2, 0)
    tensor = ((tensor +1)/2.0) * 255.0
    tensor = tensor.numpy()
    return tensor

def proprocess(image):
    
    transform = getTransform()
    trgbImage = transform(image)
    r,g,b = trgbImage[0]+1, trgbImage[1]+1, trgbImage[2]+1
    tattentionImage = 1. - ((0.299*r+0.587*g+0.114*b)/2.)
    tattentionImage = torch.unsqueeze(tattentionImage, 0)
    trgbImage = torch.unsqueeze(trgbImage, 0)
    tattentionImage = torch.unsqueeze(tattentionImage, 0)

    return trgbImage, tattentionImage

def run(inputPath, OutputPath):

    modelToLoad = torch.jit.load("./EGAN.pt")
    print("OK")
    count =0 
    for filename in os.listdir(inputPath):
        
        print("Processing Image : ", filename)
        inputImage = cv2.imread(os.path.join(inputPath,filename))
        
        rgbImage, attentionImage = proprocess(inputImage)
        
        fake, real = modelToLoad.forward(rgbImage,attentionImage )
        
        fake_B = convertToCV(fake)
        fake_B1 = cv2.cvtColor(fake_B, cv2.COLOR_RGB2BGR)
        cv2.imwrite(OutputPath + "pic1.png" , fake_B )

C++端推理代码 #

#define A 0.299
#define B 0.5870
#define C 0.114

cv::Mat torchTensortoCVMat1C(torch::Tensor& tensor)
{
    tensor = tensor.squeeze(0);
    tensor = tensor.to(torch::kCPU).to(torch::kFloat32).detach();
    tensor = tensor.permute({ 1, 2, 0 }).contiguous();
    tensor = tensor.mul(0.5).add(0.5).mul(255.0);
    tensor = tensor.to(torch::kU8);

    int64_t height = tensor.size(0);
    int64_t width  = tensor.size(1);
    cv::Mat mat    = cv::Mat(cv::Size(width, height), CV_8UC3, tensor.data_ptr<uchar>());
    return mat.clone();
}

std::vector<torch::jit::IValue> CV2Tensor(const cv::Mat& cv_Image)
{
    torch::Tensor tInputImage = (torch::from_blob(cv_Image.data, { cv_Image.rows, cv_Image.cols, cv_Image.channels() }, torch::kByte));
    tInputImage = tInputImage.to(torch::kFloat).div(255);
    tInputImage = tInputImage.sub(0.5).div(0.5).permute({ 2, 0, 1 });

    torch::Tensor red   = tInputImage[0] + 1 ;
    torch::Tensor green = tInputImage[1] + 1 ;
    torch::Tensor blue  = tInputImage[2] + 1;

    red   = red.mul(A);
    green = green.mul(B);
    blue  = blue.mul(C);

    torch::Tensor channelSum = red.add(green).add(blue);
    channelSum = channelSum.div(2.);
    torch::Tensor tGrayImage = 1. - channelSum;

    tGrayImage.unsqueeze_(0);
    tGrayImage.unsqueeze_(0);
    tInputImage.unsqueeze_(0);

    std::vector<torch::jit::IValue> input;
    input.push_back(tInputImage);
    input.push_back(tGrayImage);

    return input;
}

void enhanceImage(const std::string& Img, torch::jit::script::Module& network,const std::string& outputPath, std::string& fileName)
{
    cv::Mat rgbImage;
    cv::Mat inputImage = cv::imread(Img);
    std::vector<torch::jit::IValue> input = CV2Tensor(inputImage);

    try
    {
        auto outputs = network.forward(input).toTuple();
        torch::Tensor resultFake = outputs->elements()[0].toTensor();
    
        cv::Mat output1 = torchTensortoCVMat(resultFake);
        cv::imshow("out1.png", output1);
        cv::waitKey(0);
    }
    catch (std::exception& e)
    {
        std::cout << e.what() << std::endl;
    }
}

异常现象 #

已核对推理全流程各步骤的张量输出，确认两端张量完全一致，但最终转换得到的输出图像中，C++版本存在输入图像亮区色彩溢出的问题，效果对比如下：

• Python版本输出：
• C++版本输出：

已尝试多种方案仍未解决此问题，恳请各位提供解决思路，感谢。如有需要可以补充更多相关信息。

该问题核心是C++端张量转OpenCV Mat的逻辑和Python端不匹配，可按以下步骤修复：

1. 新增张量值截断逻辑
   Python端cv2.imwrite接收0_{255范围的浮点张量时，会自动将超出范围的值截断到0}255再转换为uint8；但libTorch中直接调用to(torch::kU8)时，超出范围的浮点值会触发取模操作（例如256转为0、257转为1），直接导致亮区色彩溢出。
   修改torchTensortoCVMat1C函数代码，在乘255后增加截断逻辑：

tensor = tensor.mul(0.5).add(0.5).mul(255.0);
// 新增截断，限制值范围在0~255
tensor = torch::clamp(tensor, 0.0, 255.0);
tensor = tensor.to(torch::kU8);

2. 对齐输出通道转换逻辑
   OpenCV的imshow、imwrite接口默认要求BGR通道顺序，而模型输出的是RGB格式，需要和Python端的cv2.COLOR_RGB2BGR逻辑对齐，在返回Mat前做通道转换：

cv::Mat mat = cv::Mat(cv::Size(width, height), CV_8UC3, tensor.data_ptr<uchar>());
// 新增RGB转BGR逻辑
cv::cvtColor(mat, mat, cv::COLOR_RGB2BGR);
return mat.clone();

3. 可选核对项
   如果修复后仍存在偏差，可核对预处理阶段的通道顺序：cv::imread默认读取的是BGR格式，若模型要求输入为RGB，需要在预处理时先做BGR转RGB操作，保证输入张量和Python端完全一致。

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