深度学习中，动态改变batchsize（自用）

小伞cos

于 2024-02-29 01:09:25 发布

阅读量295

点赞数 1

文章标签：深度学习人工智能

本文链接：https://blog.csdn.net/weixin_68494476/article/details/136359660

版权

文章介绍了如何在PyTorch中使用自定义的DynamicBatchSampler，该类允许根据需要动态调整批次大小，以提高数据加载的灵活性。作者通过实例展示了如何创建并使用这种动态批次大小的数据加载器。

摘要由CSDN通过智能技术生成

import torch
from torch.utils.data import Dataset, DataLoader, BatchSampler, SequentialSampler

# 假设有一个简单的数据集
class SimpleDataset(Dataset):
    def __init__(self, size):
        self.data = list(range(size))  # 假设数据集简单包含连续整数
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

# 自定义BatchSampler
class DynamicBatchSampler(BatchSampler):
    def __init__(self, sampler, batch_size_callback):
        self.sampler = sampler
        self.batch_size_callback = batch_size_callback  # 一个函数，用来决定批次大小
    
    def __iter__(self):
        batch = []
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size_callback():  # 调用回调函数获取当前批次大小
                yield batch
                batch = []
        if len(batch) > 0:  # 确保最后一个批次被处理
            yield batch

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

# 示例使用
dataset = SimpleDataset(100)  # 创建一个简单的数据集
a = 5  # 或者其他值，如 8

# 定义一个回调函数来决定批次大小
def get_dynamic_batch_size():
    return a

# 创建一个SequentialSampler
sampler = SequentialSampler(dataset)
# 创建DynamicBatchSampler
dynamic_batch_sampler = DynamicBatchSampler(sampler, get_dynamic_batch_size)

# 创建DataLoader
dataloader = DataLoader(dataset, batch_sampler=dynamic_batch_sampler)

# 迭代DataLoader
for batch in dataloader:
    print(batch)  # 打印每个批次

可以看出其实对于batchsize，有BatchSampler这个类，能写得非常非常灵活的