视频任务中ResNet50编码器CUDA内存不足问题求助

云原生阿岚

2026-6-19

视频任务端到端模型训练的CUDA内存不足问题

我正在基于视频任务训练一个端到端模型，采用PyTorch ResNet50作为编码器，输入形状为(1, seq_length, 3, 224, 224)，其中seq_length代表每个视频的帧数。例如视频1有1500帧时输入形状为(1,1500,3,224,224)，视频2有2000帧时为(1,2000,3,224,224)。但将输入传入ResNet时，在前向传播的第一个卷积层阶段出现CUDA内存不足的问题。

已尝试的解决方法

调用torch.cuda.empty_cache()清理显存
在DataLoader中设置pin_memory=True和prefetch_factor=2优化数据加载
设置环境变量PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True缓解显存碎片
减小seq_length，但该操作会严重破坏视频序列结构，大幅降低模型性能

求其他可行的解决方法。

完整错误信息

Traceback (most recent call last):
  File "Path/E2E.py", line 143, in <module>
    p_classes1, phase_preds = model1(long_feature)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "Path/E2E.py", line 47, in forward
    x = self.resnet_lstm(x)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "Path/train_embedding.py", line 226, in forward
    x = self.share.forward(x)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/container.py", line 219, in forward
    input = module(input)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/container.py", line 219, in forward
    input = module(input)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torchvision/models/resnet.py", line 155, in forward
    out = self.bn3(out)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "Path/lib/python3.9/site-packages/torch/nn/modules/batchnorm.py", line 176, in forward
    return F.batch_norm(
  File "Path/lib/python3.9/site-packages/torch/nn/functional.py", line 2512, in batch_norm
    return torch.batch_norm(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 614.00 MiB. GPU 0 has a total capacity of 23.65 GiB of which 353.38 MiB is free. Including non-PyTorch memory, this process has 23.28 GiB memory in use. Of the allocated memory 22.84 GiB is allocated by PyTorch, and 3.77 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

ResNet50模型代码（self.fc仅用于加载预训练模型）

class resnet_lstm(torch.nn.Module):
    def __init__(self):
        super(resnet_lstm, self).__init__()
        resnet = models.resnet50(pretrained=True)
        self.share = torch.nn.Sequential()
        self.share.add_module("conv1", resnet.conv1)
        self.share.add_module("bn1", resnet.bn1)
        self.share.add_module("relu", resnet.relu)
        self.share.add_module("maxpool", resnet.maxpool)
        self.share.add_module("layer1", resnet.layer1)
        self.share.add_module("layer2", resnet.layer2)
        self.share.add_module("layer3", resnet.layer3)
        self.share.add_module("layer4", resnet.layer4)
        self.share.add_module("avgpool", resnet.avgpool)
        self.fc = nn.Sequential(nn.Linear(2048, 512),
                                nn.ReLU(),
                                nn.Linear(512, 7))

    def forward(self, x):
        x = x.view(-1, 3, 224, 224)
        x = self.share.forward(x)
        x = x.view(1,-1, 2048)
        return x

数据集实现代码

file_paths是长度为seq_length的图像路径列表：

class CustomDataset(Dataset):
    def __init__(self, file_paths, file_labels, transform=None,
                 loader=pil_loader):
        self.file_paths = file_paths
        self.file_labels_phase = file_labels
        self.transform = transform
        self.loader = loader

    def __getitem__(self, index):
        img_names_list = self.file_paths[index]
        labels_phase = self.file_labels_phase[index]
        imgs = [self.loader(img_name) for img_name in img_names_list]
        if self.transform is not None:
            imgs = [self.transform(img) for img in imgs]

        return imgs, labels_phase, index

    def __len__(self):
        return len(self.file_paths)

class SeqSampler(RandomSampler):
    def __init__(self, data_source, seed=1):
        super().__init__(data_source)
        self.data_source = data_source
        self.seed = seed

    def __iter__(self):
        if self.seed is not None:
            random.seed(self.seed)
        
        # Generate a list of indices and shuffle them
        indices = list(range(len(self.data_source)))
        random.shuffle(indices)
        return iter(indices)

    def __len__(self):
        return len(self.idx)

train_loaders = DataLoader(
            train_dataset_80,
            batch_size=1,
            sampler=SeqSampler(train_dataset_80),
            num_workers=1,
            pin_memory=True,
        )

输入数据传入方式

for data,labels_phase in tqdm(train_loaders,desc=f"Epoch    {epoch+1}/{max_epochs}"):
        
        long_feature = torch.tensor(np.array(data)).to(device)
        labels_phase = np.asarray(labels_phase).squeeze()

        optimizer1.zero_grad()

        labels_phase = torch.LongTensor(labels_phase).to(device)
        
        #Allocated: 2239.18 MB
        #Cached:    2248.00 MB
        p_classes = model1(long_feature)

nvidia-smi输出结果

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.23.08              Driver Version: 545.23.08    CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================| |   0  NVIDIA RTX A6000               On  | 00000000:25:00.0 Off |                  Off |
| 30%   24C    P8              17W / 300W |     12MiB / 49140MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================| |    0   N/A  N/A      7388      G   /usr/lib/xorg/Xorg                            4MiB |
+---------------------------------------------------------------------------------------+

可行的解决方法

1. 分批次处理视频帧

将单视频的所有帧拆分为多个小批次处理，避免一次性加载全部帧占用过多显存，处理完成后拼接结果：

def forward(self, x):
    batch_size, seq_len, c, h, w = x.shape
    # 自定义子批次大小，根据显存调整
    chunk_size = 128
    features = []
    for i in range(0, seq_len, chunk_size):
        chunk = x[:, i:i+chunk_size, :, :, :].view(-1, c, h, w)
        feat = self.share(chunk)
        features.append(feat.view(batch_size, -1, 2048))
    x = torch.cat(features, dim=1)
    return x

2. 启用混合精度训练

通过torch.cuda.amp将部分张量从FP32转为FP16，大幅降低显存占用，同时几乎不影响模型性能：

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

# 训练循环内修改
optimizer1.zero_grad()
with autocast():
    p_classes = model1(long_feature)
    loss = ... # 替换为你的损失计算逻辑
scaler.scale(loss).backward()
scaler.step(optimizer1)
scaler.update()

3. 修改BatchNorm层行为

ResNet默认的BatchNorm会对所有视频帧计算全局均值方差，既占用显存也不符合视频时序特性。可替换为LayerNorm：

# 在resnet_lstm的__init__中添加替换逻辑
def replace_bn_with_ln(module):
    for name, child in module.named_children():
        if isinstance(child, nn.BatchNorm2d):
            setattr(module, name, nn.LayerNorm(child.num_features))
        else:
            replace_bn_with_ln(child)

replace_bn_with_ln(self.share)

4. 冻结ResNet部分层

若无需微调整个ResNet，可冻结前几层参数，减少反向传播时的显存开销：

# 在resnet_lstm的__init__中添加冻结逻辑
for param in self.share.conv1.parameters():
    param.requires_grad = False
for param in self.share.bn1.parameters():
    param.requires_grad = False
for param in self.share.layer1.parameters():
    param.requires_grad = False

可根据任务需求调整冻结层数。

5. 优化数据加载流程

避免从列表转numpy再转tensor的冗余步骤，直接在数据集内返回tensor：

# 修改CustomDataset的__getitem__方法
def __getitem__(self, index):
    img_names_list = self.file_paths[index]
    labels_phase = self.file_labels_phase[index]
    imgs = []
    for img_name in img_names_list:
        img = self.loader(img_name)
        if self.transform is not None:
            img = self.transform(img)
        imgs.append(img)
    imgs = torch.stack(imgs)
    return imgs, labels_phase, index

训练循环内直接使用：

long_feature = data.to(device)

6. 使用梯度累积

通过多次前向传播后统一更新参数，模拟大批次训练效果，同时保持单批次显存占用：

accumulation_steps = 4  # 累积4次梯度再更新参数

for step, (data,labels_phase) in enumerate(tqdm(train_loaders,desc=f"Epoch    {epoch+1}/{max_epochs}")):
    long_feature = data.to(device)
    labels_phase = torch.LongTensor(np.asarray(labels_phase).squeeze()).to(device)
    
    with autocast():
        p_classes = model1(long_feature)
        loss = ... # 替换为你的损失计算逻辑
    
    loss = loss / accumulation_steps
    scaler.scale(loss).backward()
    
    if (step + 1) % accumulation_steps == 0:
        scaler.step(optimizer1)
        scaler.update()
        optimizer1.zero_grad()

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

视频任务中ResNet50编码器CUDA内存不足问题求助

已尝试的解决方法 #

完整错误信息 #

ResNet50模型代码（self.fc仅用于加载预训练模型） #

数据集实现代码 #

输入数据传入方式 #

nvidia-smi输出结果 #

可行的解决方法 #

1. 分批次处理视频帧 #

2. 启用混合精度训练 #

3. 修改BatchNorm层行为 #

4. 冻结ResNet部分层 #

5. 优化数据加载流程 #

6. 使用梯度累积 #