如何避免Google Colab处理ISIC-2019采样时内存溢出崩溃

我使用提供83.5 GB内存的Google Colab Pro版本，针对包含17k张图像的ISIC-2019训练数据集，采用imblearn库实现随机过采样与欠采样以解决类别不平衡问题。但运行代码时RAM占用超出可用内存导致崩溃，请问是否有可避免内存溢出的类别不平衡问题替代处理方法？

原参考代码：

import torch
import torch.utils.data as data
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
from imblearn.over_sampling import RandomOverSampler
from imblearn.under_sampling import RandomUnderSampler


train_transforms = transforms.Compose([
    transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.RandomCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/train', transform=train_transforms)
test_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/test', transform=train_transforms)
val_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/val', transform=train_transforms)

# Extract the data and labels from the datasets
train_data = [np.array(sample) for sample, _ in train_dataset]
train_labels = [label for _, label in train_dataset]

# Resample the training data using random oversampling
ros = RandomOverSampler(random_state=42)
train_data_resampled, train_labels_resampled = ros.fit_resample(np.array(train_data).reshape(-1, 224 * 224 * 3), np.array(train_labels))

# Undersample the resampled training data using random undersampling
rus = RandomUnderSampler(random_state=42)
train_data_resampled, train_labels_resampled = rus.fit_resample(train_data_resampled, train_labels_resampled)

# Convert the resampled data back to PIL.Image objects and create a new dataset and dataloader
train_data_resampled = [train_data_resampled[i].reshape(3, 224, 224) for i in range(len(train_data_resampled))]
train_dataset_resampled = data.TensorDataset(torch.stack([torch.from_numpy(sample) for sample in train_data_resampled]), torch.from_numpy(train_labels_resampled))

# Create the test and train dataloader
test_dataloader = data.DataLoader(test_dataset, batch_size=32, shuffle=False)
train_dataloader = data.DataLoader(train_dataset_resampled, batch_size=32, shuffle=True)
val_dataloader = data.DataLoader(val_dataset, batch_size=32, shuffle=False)

核心问题分析 #

原代码的内存瓶颈在于：一次性将所有训练图像加载为numpy数组并做全局重采样，17k张224×224的3通道图像会生成17000×150528的超大数组（约9.5GB），过采样后数据量进一步膨胀，直接耗尽内存。下面是三种低内存占用的替代方案：

方案1：使用WeightedRandomSampler实现采样级重采样 #

无需修改原始数据集，通过DataLoader的采样器按类别权重随机采样，等效于动态过采样/欠采样，内存仅保留原始数据集。

import torch
import torch.utils.data as data
from torchvision import datasets, transforms
import numpy as np

# 定义变换
train_transforms = transforms.Compose([
    transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.RandomCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225]),
])

# 加载原始数据集
train_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/train', transform=train_transforms)
test_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/test', transform=train_transforms)
val_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/val', transform=train_transforms)

# 计算每个类别的样本数量
class_counts = np.bincount(train_dataset.targets)
# 计算类别权重：反比于样本数量
class_weights = 1.0 / class_counts
# 为每个样本分配对应类别的权重
sample_weights = class_weights[train_dataset.targets]

# 构建加权随机采样器
sampler = data.WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(train_dataset),  # 可调整为目标采样总数，比如等于多数类样本数的N倍
    replacement=True  # True表示过采样（允许重复选少数类样本），False表示欠采样
)

# 构建DataLoader，使用自定义采样器
train_dataloader = data.DataLoader(train_dataset, batch_size=32, sampler=sampler)
test_dataloader = data.DataLoader(test_dataset, batch_size=32, shuffle=False)
val_dataloader = data.DataLoader(val_dataset, batch_size=32, shuffle=False)

方案2：在线增强+类别权重损失函数 #

对少数类样本实时应用数据增强（等效在线过采样），同时在损失函数中加入类别权重，抵消类别不平衡的影响，完全避免重采样带来的内存开销。

步骤1：自定义数据集，针对少数类增强 #

from torchvision.transforms import v2

# 定义基础变换和增强变换
base_transform = transforms.Compose([
    transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.RandomCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225]),
])

# 针对少数类的增强变换
augment_transform = v2.Compose([
    base_transform,
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomVerticalFlip(p=0.5),
    v2.RandomRotation(15),
    v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1)
])

# 自定义数据集
class AugmentedDataset(data.Dataset):
    def __init__(self, dataset, minority_classes):
        self.dataset = dataset
        self.minority_classes = minority_classes
        
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, idx):
        img, label = self.dataset[idx]
        # 少数类样本应用增强
        if label in self.minority_classes:
            return augment_transform(img), label
        else:
            return base_transform(img), label

# 假设ISIC-2019中类别0、2是少数类（根据实际数据调整）
minority_classes = [0, 2]
train_dataset_aug = AugmentedDataset(train_dataset, minority_classes)

步骤2：使用带类别权重的损失函数 #

# 计算类别权重
class_counts = np.bincount(train_dataset.targets)
class_weights = torch.tensor((len(train_dataset) / (len(class_counts) * class_counts)), dtype=torch.float32)

# 定义损失函数
criterion = torch.nn.CrossEntropyLoss(weight=class_weights.to('cuda' if torch.cuda.is_available() else 'cpu'))

# 构建普通DataLoader即可
train_dataloader = data.DataLoader(train_dataset_aug, batch_size=32, shuffle=True)

方案3：分批重采样（保留imblearn逻辑但降低内存占用） #

如果一定要用imblearn的重采样方法，可分批处理数据，避免一次性加载所有图像到内存。

import torch
import torch.utils.data as data
from torchvision import datasets, transforms
import numpy as np
from imblearn.over_sampling import RandomOverSampler
from imblearn.under_sampling import RandomUnderSampler

train_transforms = transforms.Compose([
    transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.RandomCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder(root='data/isic2019/ISIC - 2019/train', transform=train_transforms)

# 分批处理参数
batch_size_resample = 1000  # 每次处理1000张图
total_samples = len(train_dataset)
resampled_data = []
resampled_labels = []

ros = RandomOverSampler(random_state=42)
rus = RandomUnderSampler(random_state=42)

for i in range(0, total_samples, batch_size_resample):
    # 提取当前批次的样本和标签
    batch_samples = train_dataset[i:i+batch_size_resample]
    batch_data = np.array([np.array(sample) for sample, _ in batch_samples])
    batch_labels = np.array([label for _, label in batch_samples])
    
    # 批次内重采样
    batch_data_reshaped = batch_data.reshape(-1, 224*224*3)
    batch_data_resampled, batch_labels_resampled = ros.fit_resample(batch_data_reshaped, batch_labels)
    batch_data_resampled, batch_labels_resampled = rus.fit_resample(batch_data_resampled, batch_labels_resampled)
    
    # 转换回张量并保存
    batch_data_resampled = batch_data_resampled.reshape(-1, 3, 224, 224)
    resampled_data.append(torch.from_numpy(batch_data_resampled))
    resampled_labels.append(torch.from_numpy(batch_labels_resampled))

# 合并所有批次
resampled_data = torch.cat(resampled_data, dim=0)
resampled_labels = torch.cat(resampled_labels, dim=0)
train_dataset_resampled = data.TensorDataset(resampled_data, resampled_labels)

# 构建DataLoader
train_dataloader = data.DataLoader(train_dataset_resampled, batch_size=32, shuffle=True)

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