如何在torchvision的ResNet50模型中使用torch.Tensor.permute适配.npy输入

核心思路 #

你需要先将.npy文件的维度[batch_size, depth, height, width, channels]调整为模型预期的[batch_size, channels, depth, height, width]，再输入模型。以下是两种实用方案：

方案1：数据预处理阶段调整（推荐） #

在读取.npy文件后，转换为PyTorch张量时直接调整维度，无需修改模型结构：

import numpy as np
import torch

# 读取npy数据
npy_data = np.load("your_input.npy")  # 形状：[batch_size, depth, H, W, channels]
# 转为PyTorch张量并调整维度顺序
input_tensor = torch.tensor(npy_data, dtype=torch.float32).permute(0, 4, 1, 2, 3)
# 此时input_tensor形状为[batch_size, channels, depth, H, W]，可直接输入模型

方案2：自定义模型内置维度调整 #

如果希望模型本身能直接接收原维度的输入，可以封装ResNet50，在forward方法中加入维度转换逻辑：

import torch.nn as nn
from torchvision.models import resnet50

class NPYCompatibleResNet50(nn.Module):
    def __init__(self, num_classes=1000):
        super().__init__()
        self.resnet = resnet50(pretrained=False, num_classes=num_classes)
        
    def forward(self, x):
        # 输入x的形状：[batch_size, depth, H, W, channels]
        x = x.permute(0, 4, 1, 2, 3)  # 转换为模型预期维度
        # 注意：torchvision原生ResNet50是2D图像模型，仅支持4维输入[batch, channels, H, W]
        # 如果你的输入带depth维度（如3D图像/视频帧），需额外处理：
        # 示例：将depth维度合并到batch，处理后聚合结果
        batch_size, channels, depth, H, W = x.shape
        x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * depth, channels, H, W)
        x = self.resnet(x)
        x = x.reshape(batch_size, depth, num_classes).mean(dim=1)  # 对depth维度结果取平均
        return x

关键注意事项 #

torchvision的ResNet50是为2D图像设计的，默认只接受4维张量。如果你的输入包含depth维度（如3D数据或视频序列），要么按上述示例拆分depth维度处理，要么改用专门的3D ResNet模型（如torchvision.models.video.r3d_50），它原生支持5维输入。

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