pytorch中dataloader中报错（RuntimeError: each element in list of batch）解决方法

2024-02-04 16:06:00 by wst

数据处理

---

       在进行数据训练的时候，首先要载入数据集。但是如果处理不好，载入数据也会报错，这里我就遇到错误：RuntimeError: each element in list of batch should be of equal size

下面给出解决方法，直接上代码：

from torch.utils.data import DataLoader
import webdataset as wds
import json


def test_dataloader():
    "test"
    def my_decoder(key, value):
        if not key.endswith(".jpg"):
            return None
        assert isinstance(value, bytes)
        return value

    url = "/home/wst/dataset/v1_res/{00000..00022}.tar"

    def mycollate(item):
        if len(item) != 128:
            return item
        print("type:", type(item))
        return len(item)
    dataset = wds.WebDataset(url).shuffle(1000).decode(my_decoder)
    dataloader = DataLoader(dataset, num_workers=4, shuffle=False, batch_size=128, collate_fn=mycollate)
    for sample in dataloader:
        print(sample)

if __name__ == "__main__":
    test_dataloader()

注意其中的mycollate函数，如果dataloader中不传入collate_fn=mycollate，则会报错，因为每个批次的数据量不一致。

这个代码如果要用，记得改改哦。

 

---