Bootstrap

torch.utils.data.dataset 的数据组织形式——python list、dict、tuple内存消耗量

在这里插入图片描述
在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.getsizeofpympler.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只能对一部分数据打乱顺序。
两种方式各有利弊。

笔者后续会单独写篇博客,来详细讲讲这两种数据读取策略。

;