从 Pytorch 数据加载器的 collat​​e_fn 获取批次

Question

我正在阅读视频并将它们分成批处理。我的批次无法容纳整个视频。所以我使用整理功能将它们分开。视频的形状为 (128,3,224,224)。 其中 128 是帧，3 是通道 (RGB)。 使用整理功能后数据加载器的输出形状为 torch.Size([4, 32, 3, 224, 224]) 其中索引 0 是批处理。我如果通过这个浴数据，如果会得到内存错误。有没有办法发一批 2 .?
这是用于重现能力的代码。
生成虚拟数据

features=[]
labels=[]
#here the data is in ram, but actually its in disk and i cant load all data to ram to manipulate. 
for i in range(5):
    features.append(np.random.randint(low=0, high=256, size=(128,3,224,224)))
    labels.append(np.random.randint(low=0, high=2))

数据集

class DataReader(Dataset):
    def __init__(self,features,labels):
        self.features = features
        self.labels = labels

    def __getitem__(self, index):
        x=self.features[index]
        y=self.labels[index]
        return x,y
        
    def __len__(self):
        return len(self.features)

数据加载器

def reshaped(i,bs=32):
    i=i[len(i)%bs::].reshape(len(i)//bs,-1,3,224,224)
    return i
def collate_fn(batch):
    data=reshaped(batch[0])
    print(data.shape)
    label=[batch[1]]*len(data)
    return data,label
dl=DataLoader(DataReader(features,labels), batch_size =None, num_workers=0,pin_memory=True,shuffle=False,collate_fn=collate_fn)
batch=next(iter(dl))

batch[0] 形状是(4, 32, 3, 224, 224) 即使是单个批次，生成的数据也很大，所以我想发送(2, 32, 3, 224, 224) 然后另一个(2, 32, 3, 224, 224)

不要与batch=None 混淆，我只是禁用自动批处理。实际上batch=None 加载一批。我在整理功能中将该批次转换为四个批次。 但问题是现在所有四个都将被进一步发送，我希望只发送其中 2 个，然后再发送 2 个到模型，而不是同时发送所有四个。

Answer 1

根据我的理解，每当您的模型转发时，都会调用 collat​​e_fn。所以，每次只取批处理中的当前元素和下一个元素。这可能不正确，您可以尝试一下吗？请告诉我结果。

def collate_fn(batch):
    data = batch[0][0: 2] # current element and next element in the list
    label = batch[1][0: 2]
    return data,label