Pytorch torch.utils.data.dataloader.default_collate 介绍

torch.utils.data.dataloader.default_collate 是 PyTorch 中 DataLoader 默认的 collate_fn 函数,用于将一个批次的样本数据合并成张量(Tensor)或其他结构化数据格式。以下是关于 default_collate 的详细介绍:

1. 功能

default_collate 的主要功能是将一个批次的样本数据(通常是列表形式)递归地打包成张量。它会根据数据的结构自动处理以下几种情况:

  • 标量:将标量打包成张量。

  • 列表或元组:将列表或元组递归打包成张量。

  • 字典:将字典的键值对分别打包成张量。

  • NumPy 数组:将 NumPy 数组转换为 PyTorch 张量。

  • 其他类型 :如果无法处理,会抛出 TypeError

2. 默认行为

以下是 default_collate 的默认行为示例:

2.1 标量

如果样本数据是标量,default_collate 会将它们打包成一个张量:

复制代码
import torch
from torch.utils.data.dataloader import default_collate

data = [1, 2, 3, 4]
batch = default_collate(data)
print(batch)  # 输出: tensor([1, 2, 3, 4])
2.2 列表或元组

如果样本数据是列表或元组,default_collate 会递归地将它们打包成张量:

复制代码
data = [[1, 2], [3, 4], [5, 6]]
batch = default_collate(data)
print(batch)  # 输出: tensor([[1, 2], [3, 4], [5, 6]])
2.3 字典

如果样本数据是字典,default_collate 会将字典的键值对分别打包成张量:

复制代码
data = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
batch = default_collate(data)
print(batch)  # 输出: {'a': tensor([1, 3, 5]), 'b': tensor([2, 4, 6])}
2.4 NumPy 数组

如果样本数据是 NumPy 数组,default_collate 会将其转换为 PyTorch 张量:

复制代码
import numpy as np

data = [np.array([1, 2]), np.array([3, 4]), np.array([5, 6])]
batch = default_collate(data)
print(batch)  # 输出: tensor([[1, 2], [3, 4], [5, 6]])

3. 局限性

虽然 default_collate 很强大,但它有一些局限性:

  • 无法处理变长序列 :如果样本数据是变长的(例如不同长度的序列),default_collate 会直接抛出错误。这种情况下需要自定义 collate_fn

  • 无法处理自定义数据格式 :如果样本数据是自定义的复杂结构(例如嵌套的字典或列表),default_collate 可能无法正确处理。

4. 自定义 collate_fn

如果 default_collate 无法满足需求,可以通过自定义 collate_fn 来实现更灵活的数据处理。例如,处理变长序列时,可以使用 torch.nn.utils.rnn.pad_sequence 来填充序列:

复制代码
import torch
from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __init__(self):
        self.data = [[1, 2], [3, 4, 5], [6], [7, 8, 9, 10]]

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

    def __getitem__(self, idx):
        return self.data[idx]

def custom_collate_fn(batch):
    sequences = [torch.tensor(seq) for seq in batch]
    padded_sequences = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True)
    return padded_sequences

dataset = MyDataset()
dataloader = DataLoader(dataset, batch_size=2, collate_fn=custom_collate_fn)

for batch in dataloader:
    print(batch)
    # 输出:
    # tensor([[1, 2, 0],
    #         [3, 4, 5]])
    # tensor([[6, 0, 0],
    #         [7, 8, 9]])

5. 总结

  • default_collate 是 PyTorch 中 DataLoader 的默认 collate_fn,用于将样本数据打包成张量。

  • 它可以处理标量、列表、元组、字典和 NumPy 数组等数据类型。

  • 如果数据具有特殊结构(如变长序列或自定义格式),需要自定义 collate_fn 来灵活处理。

通过理解 default_collate 的行为,可以更好地决定是否需要自定义 collate_fn 来满足特定需求。

相关推荐
测试员周周5 小时前
【Appium 系列】第16节-WebView-H5上下文切换 — 混合应用的自动化难点
运维·开发语言·人工智能·功能测试·appium·自动化·测试用例
测试19985 小时前
软件测试 - 单元测试总结
自动化测试·软件测试·python·测试工具·职场和发展·单元测试·测试用例
K姐研究社7 小时前
怎么用AI制作电商口播视频,开拍APP一键生成
人工智能·音视频
LaughingZhu8 小时前
Product Hunt 每日热榜 | 2026-05-21
前端·人工智能·经验分享·chatgpt·html
曲幽8 小时前
我用了FastApiAdmin后,连夜把踩过的坑都整理出来了
redis·python·postgresql·vue3·fastapi·web·sqlalchemy·admin·fastapiadmin
传说故事8 小时前
【论文阅读】MotuBrain: An Advanced World Action Model for Robot Control
论文阅读·人工智能·具身智能·wam
北京耐用通信9 小时前
全域适配工业场景耐达讯自动化Modbus TCP 转 PROFIBUS 网关轻松实现以太网与现场总线互通
网络·人工智能·网络协议·自动化·信息与通信
火山引擎开发者社区9 小时前
TRAE × 火山引擎 Supabase:为你的 AI 应用装上“数据引擎”
人工智能
小a彤9 小时前
GE 在 CANN 五层架构中的位置
人工智能·深度学习·transformer
前端若水9 小时前
会话管理:创建、切换、删除对话历史
前端·人工智能·python·react.js