强化学习视频分类训练时CUDA内存不足问题求助

问题背景 #

使用以下代码训练视频分类模型（注：代码实际为CNN结构，非强化学习），输入数据维度为129x15x3x60x230x230，单视频片段张量维度为3x60x230x230（3通道、60帧、230×230分辨率）。在RTX 4070Ti（12GB显存）上运行时触发内存不足错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 48.43 GiB. GPU

已尝试调整batch size、逐帧输入等操作，均无效果，且不愿通过缩小帧尺寸丢失数据，寻求可行优化方案。

用户代码 #

import os
import gym
from gym import spaces
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from tqdm import tqdm
from TTset_main import load_TTset

# Load the dataset
filted_tset_path = "training_path.csv"
filted_vset_path = "validation_path.csv"
TTset_training, TTset_vali, training_data_length = load_TTset(filted_tset_path, filted_vset_path)

# Custom Dataset Class
class VideoDataset(Dataset):
    def __init__(self, data_list):
        self.data_list = data_list
        self.length = sum(len(video[0]) for video in data_list)  # Total number of frames

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        video_idx = 0
        frame_idx = idx
        # Find the corresponding video and frame
        for video, labels in self.data_list:
            if frame_idx < len(video):
                frame = video[frame_idx]
                label = labels[frame_idx]
                return frame, label
            frame_idx -= len(video)
        raise IndexError("Index out of range")

# Create DataLoader
batch_size = 1  # One frame at a time
video_dataset = VideoDataset(TTset_training)
data_loader = DataLoader(video_dataset, batch_size=batch_size, shuffle=True)

# CNN Model
class CNNModel(nn.Module):
    def __init__(self, num_classes):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv3d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv3d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(64*30*115*115, 512)
        self.fc2 = nn.Linear(512, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool3d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool3d(x, 2)
        x = x.view(-1, 64*30*115*115)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Training Loop
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CNNModel(num_classes=20).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_function = nn.CrossEntropyLoss()

# Training with progress bar
total_epochs = 10
for epoch in range(total_epochs):
    training_acu_count = 0  # Reset the accuracy count for each epoch
    for i, (x, y) in tqdm(enumerate(data_loader)):
        model.train()
        x = x.to(device).unsqueeze(0)  # Add batch dimension
        y = y.to(device)
        pred = model(x)
        t_loss = loss_function(pred, y.unsqueeze(0))  # Match dimensions for loss calculation
        t_loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
        pred_convert = torch.argmax(pred, 1)
        training_acu_count += (pred_convert == y).sum().item()

    print(f"Epoch {epoch+1}/{total_epochs}, Training Accuracy: {training_acu_count/len(video_dataset):.4f}")

# Save the model
torch.save(model.state_dict(), "video_classification_model.pth")

可行优化方案 #

1. 修正核心维度错误（解决OOM根源） #

你的代码存在两个致命维度错误，直接导致显存爆炸：

• 全连接层维度计算错误：两次max_pool3d(stride=2)后，时间维度60→30→15，空间维度230→115→57，原代码错误使用64*30*115*115作为全连接层输入，导致该层参数高达2500万，仅这一层就占用约10GB显存。
• 多余维度冗余：DataLoader的batch_size=1已为输入添加batch维度，额外的unsqueeze(0)会把输入变成1x1x3x60x230x230，大幅增加卷积计算量。

修正后的模型：

class CNNModel(nn.Module):
    def __init__(self, num_classes):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv3d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv3d(32, 64, kernel_size=3, stride=1, padding=1)
        self.global_avg_pool = nn.AdaptiveAvgPool3d(1)  # 全局池化替代手动flatten，大幅减少参数
        self.fc1 = nn.Linear(64, 512)
        self.fc2 = nn.Linear(512, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool3d(x, 2)  # 输出: 32x30x115x115
        x = F.relu(self.conv2(x))
        x = F.max_pool3d(x, 2)  # 输出: 64x15x57x57
        x = self.global_avg_pool(x)  # 输出: 64x1x1x1
        x = x.view(-1, 64)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

2. 开启混合精度训练 #

使用PyTorch自动混合精度，在几乎不损失精度的前提下，将显存占用降低约50%：

from torch.cuda.amp import GradScaler, autocast

# 初始化混合精度组件
scaler = GradScaler()

# 修改训练循环
for epoch in range(total_epochs):
    training_acu_count = 0
    for i, (x, y) in tqdm(enumerate(data_loader)):
        model.train()
        x = x.to(device)  # 移除多余的unsqueeze(0)
        y = y.to(device)
        
        with autocast():
            pred = model(x)
            t_loss = loss_function(pred, y)  # 移除y的unsqueeze(0)，batch维度已匹配
        
        scaler.scale(t_loss).backward()
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
        
        pred_convert = torch.argmax(pred, 1)
        training_acu_count += (pred_convert == y).sum().item()

3. 梯度累积模拟大batch #

如果显存仍紧张，可通过梯度累积替代大batch，既保证训练效果又不增加单步显存占用：

accumulation_steps = 8  # 累积8步后更新一次参数

for epoch in range(total_epochs):
    training_acu_count = 0
    for i, (x, y) in tqdm(enumerate(data_loader)):
        model.train()
        x = x.to(device)
        y = y.to(device)
        
        with autocast():
            pred = model(x)
            t_loss = loss_function(pred, y)
            t_loss = t_loss / accumulation_steps  # 损失归一化
        
        scaler.scale(t_loss).backward()
        
        # 达到累积步数后更新参数
        if (i + 1) % accumulation_steps == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
        
        pred_convert = torch.argmax(pred, 1)
        training_acu_count += (pred_convert == y).sum().item()

4. 模型结构轻量化 #

用深度可分离卷积替换标准3D卷积，大幅减少卷积层参数：

class DepthwiseSeparableConv3d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super().__init__()
        # 深度卷积：每个通道单独卷积
        self.depthwise = nn.Conv3d(in_channels, in_channels, kernel_size, stride, padding, groups=in_channels)
        # 点卷积：融合通道信息
        self.pointwise = nn.Conv3d(in_channels, out_channels, kernel_size=1)
    
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

# 修改模型卷积层
class CNNModel(nn.Module):
    def __init__(self, num_classes):
        super(CNNModel, self).__init__()
        self.conv1 = DepthwiseSeparableConv3d(3, 32)
        self.conv2 = DepthwiseSeparableConv3d(32, 64)
        self.global_avg_pool = nn.AdaptiveAvgPool3d(1)
        self.fc1 = nn.Linear(64, 512)
        self.fc2 = nn.Linear(512, num_classes)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool3d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool3d(x, 2)
        x = self.global_avg_pool(x)
        x = x.view(-1, 64)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

5. 数据加载优化 #

• 确保load_TTset仅在CPU存储数据，训练时逐帧移到GPU，避免一次性加载所有数据到显存；
• 开启DataLoader的pin_memory加速CPU到GPU的数据传输：

data_loader = DataLoader(video_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)

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