PyTorch中Torch.utils.data的DataLoader加载数据时batch_size变了

在pytorch训练数据,发现迭代到某一个次数时,就会报错,大概意思是输入的数据的batch size变了,不是预设置的了,后来发现是在DataLoader中有一个参数,控制dataset中的数据个数不是batch_size的整数倍时,剩下的不足batch size个的数据是否会被丢弃。

DataLoader的函数定义如下:

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, 
num_workers=0, collate_fn=default_collate, pin_memory=False, 
drop_last=False)

drop_last:dataset中的数据个数可能不是batch_size的整数倍,drop_last为True会将多出来不足一个batch的数据丢弃。

所以就在代码里加上了这个参数为True,继续训练就不再报错了。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章