Pytorch中的torch.utils.data.random_split函数

📚博客主页：knighthood2001
✨公众号：认知up吧 （目前正在带领大家一起提升认知，感兴趣可以来围观一下）
🎃知识星球：【认知up吧|成长|副业】介绍
❤️如遇文章付费，可先看看我公众号中是否发布免费文章❤️
🙏笔者水平有限，欢迎各位大佬指点，相互学习进步！

在 PyTorch 中， torch.utils.data.random_split，它用于随机分割一个数据集成为训练集和验证集（或测试集）。这个函数通常用于处理 PyTorch 中的数据加载器（torch.utils.data.Dataset），例如 torchvision.datasets 中的数据集或自定义的数据集。

参数

torch.utils.data.random_split(dataset, lengths, generator=None)

• dataset: 要分割的数据集，通常是一个 torch.utils.data.Dataset 类型的对象。
• lengths: 一个整数列表，指定每个分割后的数据集的长度。例如，[train_length, val_length] 表示分割成训练集和验证集，各自的长度。
• generator: 可选参数，用于控制数据集的随机性。如果为 None，将使用默认的随机数生成器。

返回值

返回值是一个包含分割后数据集的列表，每个元素都是 Subset 对象，可以像普通数据集一样用于数据加载器。

示例用法

import torch
from torch.utils.data import Dataset, DataLoader, random_split

# 假设有一个自定义的数据集 MyDataset
class MyDataset(Dataset):
    def __init__(self):
        self.data = list(range(100))  # 假设数据集有100个样本
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        return self.data[index]

# 创建数据集对象
dataset = MyDataset()

# 定义分割长度，假设将数据集分为80个训练样本和20个验证样本
train_length = 80
val_length = 20

# 使用 random_split 函数进行分割
train_dataset, val_dataset = random_split(dataset, [train_length, val_length])

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)

# 使用示例
for batch in train_loader:
    # 训练集批次处理
    pass

for batch in val_loader:
    # 验证集批次处理
    pass

注意事项

• 数据集长度: 分割时，确保分割长度的总和等于数据集的长度。
• 随机性: random_split 函数根据指定的随机数生成器或默认生成器进行随机分割，确保数据在不同集合间的随机性。
• 灵活性: 可以根据具体需求，将数据集分成多个部分，如训练集、验证集和测试集，或根据交叉验证的需要进行分割。

使用 torch.utils.data.random_split 函数能够方便地处理数据集的划分工作，适用于多种机器学习任务中。