iOS端PyTorch模型输出at::Tensor转换为UIImage的方法咨询

你可以直接通过 PyTorch 提供的张量操作接口，把输出张量处理为符合要求的unsigned char*数据，再调用你现有的转换方法即可，具体步骤如下：

前提说明 #

默认你的输出tensor_img是PyTorch常规的NCHW格式（维度为[1, 3, 图像高度, 图像宽度]），数值范围为0~1的浮点类型，如果你的输出格式不同，可以对应调整后续处理逻辑。

实现代码 #

注意这段代码需要写在后缀为.mm的Objective-C++文件中，才能同时兼容C++的PyTorch接口和Objective-C语法：

// 推理代码执行完成后，接下面的处理逻辑
int img_height = tensor_img.size(2);
int img_width = tensor_img.size(3);

// 第一步：张量格式转换
torch::Tensor processed_tensor = tensor_img
    .to(torch::kCPU) // 确保张量转移到CPU内存，不能直接读取GPU上的张量数据
    .squeeze(0) // 去掉第0位的batch维度，维度变为[3, H, W]
    .permute({1, 2, 0}) // 维度重排为HWC格式 [H, W, 3]，适配RGB像素连续排列的要求
    .contiguous() // 确保张量内存是连续的，否则读取到的像素数据会错乱
    .mul(255) // 把0~1的浮点数值放大到0~255区间
    .clamp(0, 255) // 截断超出0-255范围的异常值
    .to(torch::kUInt8); // 转换为无符号8位整型，对应unsigned char的类型

// 第二步：获取unsigned char*数据指针
unsigned char* rgb_buffer = (unsigned char*)processed_tensor.data_ptr();

// 第三步：调用你已有的转换方法得到UIImage
UIImage *resultImage = [self convertRGBBufferToUIImage:rgb_buffer withWidth:img_width withHeight:img_height];

常见适配调整 #

• 如果你的张量输出数值范围是-1_{1，只需要在`.mul(255)`前加一行`.add(1).div(2)`，把数值映射到0}1区间即可。
• 如果输出是单通道灰度图，需要把单通道数据复制为R、G、B三个通道数值相同的三通道数据，再调用转换方法，或者单独修改转换函数适配单通道输入。
• 不需要手动释放rgb_buffer的内存，PyTorch会自动管理processed_tensor的生命周期，只要在processed_tensor被销毁前完成转换即可。

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