torch.Tensor与float指针互转后EXR图像异常问题求助

问题描述 #

我有一个形状为(512, 512, 3)的torch.Tensor类型图像x（注：描述中“首维度为颜色通道数”存在矛盾，实际形状首维度是512，通道数为第三个维度），使用imageio.imsave保存为.exr文件输出正常。需要将x的指针传给C++库，通过以下代码获取指针：

p = ctypes.cast(x.data_ptr(), ctypes.POINTER(ctypes.c_float))

再用自定义的as_tensor函数将指针转回张量：

def as_tensor(pointer, shape, torch_type):
    return torch.frombuffer((pointer._type_ * prod(shape)).from_address(ctypes.addressof(pointer.contents)), dtype = torch_type).view(*shape)

执行y = as_tensor(p, x.shape, torch.float)后，保存y为.exr得到的图像是原图像的3x3重复网格（每行对应一个通道），期望x与y数据完全一致，请问问题出在哪里？

核心原因 #

1. 内存布局不匹配：PyTorch张量默认采用C风格（行优先）的连续内存布局，view方法仅在张量内存连续且步长符合逻辑时才能正确塑形。你的as_tensor函数直接从指针创建一维张量后强制view，没有考虑原张量的步长（stride），导致数据读取顺序错误。
2. 形状描述或使用错误：如果原图像实际是通道首的格式（形状应为(3,512,512)），但你错误地用(512,512,3)去塑形，会把单个通道的空间数据当成多通道像素数据，直接导致图像出现重复网格。

解决办法 #

方法1：匹配原张量的步长（推荐） #

如果原张量x是连续的（可通过x.is_contiguous()验证），修改as_tensor函数，通过torch.as_strided精确控制张量的形状和步长，确保和原张量一致：

from functools import reduce
import operator
import ctypes
import torch

def prod(iterable):
    return reduce(operator.mul, iterable, 1)

def as_tensor(pointer, shape, torch_type, stride):
    num_elements = prod(shape)
    # 从指针创建一维张量
    flat_tensor = torch.frombuffer(
        (pointer._type_ * num_elements).from_address(ctypes.addressof(pointer.contents)),
        dtype=torch_type
    )
    # 用原张量的步长恢复正确的张量结构
    return torch.as_strided(flat_tensor, size=shape, stride=stride)

调用时传入原张量的步长：

y = as_tensor(p, x.shape, torch.float, stride=x.stride())

方法2：强制调整维度顺序（适用于形状描述错误的情况） #

如果原张量实际是通道首格式（形状(3,512,512)），但你之前误写为(512,512,3)，恢复后需要转置维度：

# 假设原张量正确形状是(3,512,512)，恢复时先按正确形状塑形
y = as_tensor(p, (3,512,512), torch.float)
# 如果需要转为(512,512,3)格式用于保存
y = y.permute(1,2,0)

方法3：确保张量连续后塑形 #

如果不确定原张量的步长，可先将一维张量转为连续张量，再调整维度：

def as_tensor(pointer, shape, torch_type):
    num_elements = prod(shape)
    flat_tensor = torch.frombuffer(
        (pointer._type_ * num_elements).from_address(ctypes.addressof(pointer.contents)),
        dtype=torch_type
    )
    # 转为连续张量后再塑形，避免步长错误
    return flat_tensor.view(*shape).contiguous()

验证步骤 #

1. 先确认原张量的基本信息：

print("原张量形状:", x.shape)
print("原张量步长:", x.stride())
print("是否连续:", x.is_contiguous())

2. 对比x和y的数值是否一致：

print("张量是否相等:", torch.allclose(x, y))

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