如何在PyTorch中向预训练ResNet50模型传输训练数据？

你当前的实现是将全量数据一次性读取到内存后再转张量，数据集规模较大时必然会出现耗时长、内存占用过高的问题。PyTorch官方提供的Dataset+DataLoader流式加载方案是更优的选择，通过懒加载+多进程加载的机制完美解决上述问题。

实现步骤 #

• 第一步：自定义数据集类，继承torch.utils.data.Dataset，实现核心的长度查询和单样本读取逻辑
• 第二步：用DataLoader封装数据集，自动完成批处理、乱序、多进程加速等操作

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

# 自定义数据集类
class ImageDataset(Dataset):
    def __init__(self, df):
        self.df = df
    
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        # 单次只读取一个样本，不会全量加载
        img = self.df['Image'].iloc[idx] 
        label = self.df['Label'].iloc[idx]
        # 转张量+调整通道顺序，单样本处理耗时可以忽略
        img_tensor = torch.tensor(img).permute(2, 0, 1) # 转成(4,512,512)
        label_tensor = torch.tensor(label).long()
        return img_tensor, label_tensor

# 实例化数据集和加载器
train_dataset = ImageDataset(train_df)
# batch_size根据显存大小调整，num_workers根据CPU核心数调整，一般设为CPU核心数的一半即可
train_loader = DataLoader(
    train_dataset,
    batch_size=8,
    shuffle=True,
    num_workers=4,
    pin_memory=True # 用GPU训练时开启，能加快数据传输到显存的速度
)

# 训练时直接遍历train_loader即可
for imgs, labels in train_loader:
    # 此处imgs维度为(batch_size,4,512,512)，labels维度为(batch_size,)，可直接传入模型
    outputs = model(imgs)
    # 后续训练逻辑

额外注意事项 #

• 官方预训练的ResNet50默认输入为3通道图像，使用4通道输入时需要修改模型第一层卷积的输入通道数：

from torchvision import models
model = models.resnet50(pretrained=True)
# 修改第一个卷积层的输入通道为4
model.conv1 = torch.nn.Conv2d(4, 64, kernel_size=(7,7), stride=(2,2), padding=(3,3), bias=False)

• 如果数据集规模很小，也可以使用全量加载的优化版本，比原始循环实现效率高很多：

# 直接批量转数组，避免循环append的开销
x_train = torch.tensor(np.stack(train_df['Image'].values)).permute(0,3,1,2)
y_train = torch.tensor(train_df['Label'].values).view(-1)

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