如何将 3D 火炬张量切成 2D 切片答案

【问题标题】：How to slice 3D torch tensor into 2D slices如何将 3D 火炬张量切成 2D 切片
【发布时间】：2020-11-08 02:54:34
【问题描述】：

我正在处理 3D CT 医学数据，并尝试将其分割成 2D 切片，以便输入到 UNet 模型中。

我已将数据加载到 Torch 数据加载器中，当前每次迭代都会产生一个 4D 张量：

for batch_index, batch_samples in enumerate(train_loader):
    data, target = batch_samples['image'].float().cuda(), batch_samples['label'].float().cuda()
    print(data.size())
torch.Size([1, 333, 512, 512])
torch.Size([1, 356, 512, 512])

比如这个。我想遍历 333 个切片，然后是 356 个切片，这样模型每次都会收到火炬大小 [1, 1, 512, 512]。

我希望是这样的：

for x in (data[:,x,:,:]):

可以，但它说我需要先定义 x。如何迭代 Torch 张量中的特定维度？

【问题讨论】：

标签： iteration pytorch torch medical dataloader

【解决方案1】：

只需指定尺寸：

for i in range(data.shape[1]):  # dim=1
    x = data[:, i, :, :]
    # [...]

如果你真的需要那个额外的维度，只需添加.unsqueeze()：

d = 1
for i in range(data.shape[d]):         # dim=1
    x = data[:, i, :, :].unsqueeze(d)  # same dim=1
    # [...]

【讨论】：