Pytorch中collate_fn()

raimm

于 2024-06-19 09:12:04 发布

阅读量223

点赞数 7

文章标签： pytorch 人工智能 python 深度学习算法

本文链接：https://blog.csdn.net/qq_52298134/article/details/139789861

版权

在 PyTorch 的 DataLoader 中，collate_fn 是一个用于将单个样本组合成批次的函数。默认情况下，DataLoader 会将来自 Dataset 的单个样本打包成批次（batch），但是对于一些复杂的情况，你可能需要自定义这个打包过程。collate_fn 就是为此目的而设计的。

`collate_fn` 的作用

批处理的灵活性：你可以通过 collate_fn 来控制如何将多个样本组合成一个批次。
处理不规则数据：对于那些形状不规则或大小不同的数据，例如文本数据（不同长度的句子），你可以自定义 collate_fn 来进行适当的填充（padding）。
数据增强和预处理：你可以在 collate_fn 中进行数据增强或其他预处理操作。

默认行为

默认情况下，DataLoader 会将 Dataset 中的样本按照批次大小自动组合，并将每个批次的样本堆叠到一起。对于大多数常见的数据类型（例如张量），这是有效的。然而，对于一些复杂的数据结构，默认行为可能不够用，这时你就需要自定义 collate_fn。

自定义 `collate_fn`

以下是一个简单的示例，演示如何自定义 collate_fn 函数：

示例 1：处理变长序列（例如文本数据）

假设我们有一个包含不同长度序列的数据集，我们希望在将这些序列组合成批次时对其进行填充（padding），使得批次中的每个序列长度相同。

import torch
from torch.utils.data import DataLoader

# 自定义 Dataset
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

# 自定义 collate_fn
def collate_fn(batch):
    # 找到批次中最长的序列
    max_length = max(len(x) for x in batch)
    
    # 对每个序列进行填充，使其长度相同
    padded_batch = []
    for seq in batch:
        padded_seq = seq + [0] * (max_length - len(seq))  # 假设 0 是填充值
        padded_batch.append(padded_seq)
    
    # 转换为张量
    return torch.tensor(padded_batch)

# 创建数据集和 DataLoader
data = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
dataset = MyDataset(data)
dataloader = DataLoader(dataset, batch_size=2, collate_fn=collate_fn)

# 迭代 DataLoader
for batch in dataloader:
    print(batch)

输出：

tensor([[1, 2, 3, 0],
        [4, 5, 0, 0]])
tensor([[6, 7, 8, 9]])

示例 2：处理复杂的数据结构

假设我们的数据集中的每个样本是一个包含多个张量和其他类型数据的字典，我们希望自定义 collate_fn 来组合这些复杂的数据结构。

import torch
from torch.utils.data import DataLoader

# 自定义 Dataset
class MyComplexDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

# 自定义 collate_fn
def collate_fn(batch):
    batch_size = len(batch)
    
    # 初始化存储批次数据的容器
    batched_data = {
        'features': torch.zeros(batch_size, 3),  # 假设每个样本的 'features' 是一个形状为 (3,) 的张量
        'labels': torch.zeros(batch_size),       # 假设每个样本的 'labels' 是一个标量
        'metadata': []                           # 假设 'metadata' 是其他类型的数据
    }
    
    # 将每个样本的数据填充到容器中
    for i, sample in enumerate(batch):
        batched_data['features'][i] = torch.tensor(sample['features'])
        batched_data['labels'][i] = torch.tensor(sample['labels'])
        batched_data['metadata'].append(sample['metadata'])
    
    return batched_data

# 创建数据集和 DataLoader
data = [
    {'features': [1.0, 2.0, 3.0], 'labels': 0, 'metadata': 'sample1'},
    {'features': [4.0, 5.0, 6.0], 'labels': 1, 'metadata': 'sample2'},
    {'features': [7.0, 8.0, 9.0], 'labels': 0, 'metadata': 'sample3'}
]
dataset = MyComplexDataset(data)
dataloader = DataLoader(dataset, batch_size=2, collate_fn=collate_fn)

# 迭代 DataLoader
for batch in dataloader:
    print(batch)

输出：

{
    'features': tensor([[1., 2., 3.],
                        [4., 5., 6.]]),
    'labels': tensor([0., 1.]),
    'metadata': ['sample1', 'sample2']
}
{
    'features': tensor([[7., 8., 9.]]),
    'labels': tensor([0.]),
    'metadata': ['sample3']
}