在Pytorch中,我们需要通过torch.utils.data.dataset
来实现数据的读取。torch.utils.data.dataset
是一种非流式的数据读取策略,需要将数据一次性导入至内存中.如果数据规模过大,可能存在内存不够的问题。
import torch
from torch.utils.data import Dataset
class TestDataset(Dataset):
def __init__(self, image_path, label_path):
'''
从image_path中读取图片
代码实现xxx
'''
'''
从label_path中读取标签
代码实现xxx
'''
'''
保存为data,每一个元素为data和target
data = ...
'''
self.data = data
def __getitem__(self, index):
'''
获取图像tensor和标签label
代码实现xxx
'''
return image_tensor, label
def __len__(self):
return len(self.data)
例如代码中,将所有image_path和label都写到self.data
中,可能会爆内存。如何优化呢?
优化
优化的关键点在于怎么减少self.data
的内存占用,其中每一个元素可以采用list、dict或tuple这三种python产用的数据结构,哪一种在内存效率上最优呢?
答案是:
tuple < list < dict
实验
有两种统计python对象的内存占用大小,分别是sys.getsizeof
和pympler.asizeof
sys.getsizeof
只返回对象本身占用的内存大小,不包括对象引用的其他对象的内存大小pympler.asizeof
返回对象及其所有引用的对象的总内存大小
所以采用pympler.asizeof
来统计对象内存占用
from pympler import asizeof
tuple_info = [('pth_{}'.format(i),'label_{}'.format(i)) for i in range(1000000)]
print("tuple size: {} bytes".format(asizeof.asizeof(tuple_info)))
## tuple size: 192440648 bytes
list_info = [['pth_{}'.format(i),'label_{}'.format(i)] for i in range(1000000)]
print("list size: {} bytes".format(asizeof.asizeof(list_info)))
## list size: 208440648 bytes
dict_info = [{'pth': 'pth_{}'.format(i),'label': 'label_{}'.format(i)} for i in range(1000000)]
print("dict size: {} bytes".format(asizeof.asizeof(dict_info)))
## dict size: 368440760 bytes
可以发现,在1000000条数据情况下,利用tuple存储比利用dict存储减少了47%的内存占用,可以省下多少钱啊!!!!!!
原理
- python中的dict涉及到哈希表的实现,因此需要额外的存储空间。不仅如此,python3.7+的dict是有序dict,也需要下标数组来维持有序特性,需要对下标数据做存储,因此,dict的存储效率最低。
- python中的list是动态数组,需要额外的内存来存储引用,用于适应list内数据调整。list的每个元素都有一个指向实际数据的引用指针,这会导致额外的内存开销。
- python中的tuple是不可变的,无额外的引用指针,导致存储效率更高。
如果真的涉及特别大规模的数据读取,建议采用torch.utils.data.IterableDataset
,为流式的数据读取策略,每次仅读取一部分数据进到内存中,一定不会爆内存。但存在的问题是随机性不够高。
torch.utils.data.dataset
预先知道了所有数据,支持对所有数据打乱顺序。
torch.utils.data.IterableDataset
只能对一部分数据打乱顺序。
两种方式各有利弊。
笔者后续会单独写篇博客,来详细讲讲这两种数据读取策略。