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

在Pytorch中，我们需要通过torch.utils.data.dataset来实现数据的读取。torch.utils.data.dataset是一种非流式的数据读取策略，需要将数据一次性导入至内存中.如果数据规模过大，可能存在内存不够的问题。

python 复制代码

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来统计对象内存占用

python 复制代码

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只能对一部分数据打乱顺序。

两种方式各有利弊。

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