YOLOX源码之【数据缓存】

这里首先需要了解下装饰器 - 廖雪峰的官方网站的用法，后面会用到。

如果cache=True，在launch前就调用get_dataset，否则launch后再调用get_dataset。

函数get_dataset调用COCODataset类，并赋给self.dataset。COCODataset继承自CacheDataset，CacheDataset继承自Dataset，Dataset继承自torch.utils.data.data.Dataset。

COCODataset在__init__中super().__init__()初始化父类CacheDataset，CacheDataset在__init__中调用self.cache_images进行缓存图片操作，代码如下。

def cache_images(
    self,
    num_imgs=None,
    data_dir=None,
    cache_dir_name=None,
    path_filename=None,
):
    assert num_imgs is not None, "num_imgs must be specified as the size of the dataset"
    if self.cache_type == "disk":
        assert (data_dir and cache_dir_name and path_filename) is not None, \
            "data_dir, cache_name and path_filename must be specified if cache_type is disk"
        self.path_filename = path_filename

    mem = psutil.virtual_memory()  # 获取系统虚拟内存信息
    mem_required = self.cal_cache_occupy(num_imgs)
    gb = 1 << 30  # 1 << 30 == 2^30 == (2^10)^3 == 1024^3
    # 1 << 30将二进制数1左移30位，其余位都为0。2^30的二进制表示是在最高位为1，其余位都为0的二进制数，即10后面跟着30个0。因此1 << 30 == 2^30

    if self.cache_type == "ram":
        if mem_required > mem.available:
            self.cache = False
        else:
            logger.info(
                f"{mem_required / gb:.1f}GB RAM required, "
                f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB RAM available, "
                f"Since the first thing we do is cache, "
                f"there is no guarantee that the remaining memory space is sufficient"
            )

    if self.cache and self.imgs is None:
        if self.cache_type == 'ram':
            self.imgs = [None] * num_imgs
            logger.info("You are using cached images in RAM to accelerate training!")
        else:   # 'disk'
            if not os.path.exists(self.cache_dir):
                os.mkdir(self.cache_dir)
                logger.warning(
                    f"\n*******************************************************************\n"
                    f"You are using cached images in DISK to accelerate training.\n"
                    f"This requires large DISK space.\n"
                    f"Make sure you have {mem_required / gb:.1f} "
                    f"available DISK space for training your dataset.\n"
                    f"*******************************************************************\\n"
                )
            else:
                logger.info(f"Found disk cache at {self.cache_dir}")
                return

        logger.info(
            "Caching images...\n"
            "This might take some time for your dataset"
        )

        num_threads = min(8, max(1, os.cpu_count() - 1))
        b = 0
        load_imgs = ThreadPool(num_threads).imap(
            partial(self.read_img, use_cache=False),  # 偏函数，固定参数use_cache=False
            # 这里是partial的一个神奇用法，修改装饰器的参数
            range(num_imgs)
        )  # 这里load_imgs是一个迭代器
        pbar = tqdm(enumerate(load_imgs), total=num_imgs)
        for i, x in pbar:   # x = self.read_img(self, i, use_cache=False)
            if self.cache_type == 'ram':
                self.imgs[i] = x
            else:   # 'disk'
                cache_filename = f'{self.path_filename[i].split(".")[0]}.npy'
                cache_path_filename = os.path.join(self.cache_dir, cache_filename)
                os.makedirs(os.path.dirname(cache_path_filename), exist_ok=True)
                np.save(cache_path_filename, x)
            b += x.nbytes
            pbar.desc = \
                f'Caching images ({b / gb:.1f}/{mem_required / gb:.1f}GB {self.cache_type})'
        pbar.close()

缓存有ram和disk两种类型，ram是一次性将训练集中所有图片读取完放到一个列表中赋给self.imgs，disk是读取每张图片并以.npy格式保存到硬盘中。首先通过psutil.virtual_memory()获取系统虚拟内存信息，然后调用self.cal_cache_occupy()计算训练集中所有图片占用内存大小。然后用多线程的方式读取图片。

这里需要特别介绍一下读取图片的操作。首先functools.partial的作用是在原始函数的基础上固定某些参数创建一个新的可调用对象，这个新的可调用对象可以像原始函数一样被调用，但是某些参数已经被预先设置好了。下面是一个例子，在这个例子中，partial(add, 5)创建了一个新的函数add_five，它实际上是add函数的一个版本，只不过把第一个参数固定为5。这样当我们调用add_five(3)时，实际上是调用add(5, 3)，所以结果是8。

from functools import partial

def add(x, y):
    return x + y

# 使用partial固定第一个参数
add_five = partial(add, 5)

print(add_five(3))  # 输出 8

partial(self.read_img, use_cache=False)调用的self.read_img在COCODataset中实现，并且将参数use_cache固定为False，但是我们看到函数read_img并没有入参use_cache，而装饰器@cache_read_img有入参use_cache，这里是partial的一个特别的用法，即可以改变装饰器的自身的参数。

@cache_read_img(use_cache=True)
# 实际调用的是cache_read_img(use_cache=True)(read_img)(self, index)
def read_img(self, index):
    return self.load_resized_img(index)

装饰器cache_read_img的实现如下，可以看到当use_cache=True时根据缓存类型从ram或disk中读取图片，当use_cache=False时调用被装饰函数read_img读取图片。这里本身就是在进行缓存图片的操作，图片还没缓存呢当然就不能从缓存中读取图片了。 

def cache_read_img(use_cache=True):
    def decorator(read_img_fn):
        """
        Decorate the read_img function to cache the image

        Args:
            read_img_fn: read_img function
            use_cache (bool, optional): For the decorated read_img function,
                whether to read the image from cache.
                Defaults to True.
        """
        @wraps(read_img_fn)  # 保持被装饰函数read_img_fn的__name__属性不变
        def wrapper(self, index, use_cache=use_cache):
            cache = self.cache and use_cache
            if cache:
                if self.cache_type == "ram":
                    img = self.imgs[index]
                    img = copy.deepcopy(img)
                elif self.cache_type == "disk":
                    img = np.load(
                        os.path.join(
                            self.cache_dir, f"{self.path_filename[index].split('.')[0]}.npy"))
                else:
                    raise ValueError(f"Unknown cache type: {self.cache_type}")
            else:
                img = read_img_fn(self, index)
            return img
        return wrapper
    return decorator

至此就完成了缓存图片的操作。