FFCV:让数据加载不再是训练模型的瓶颈

GiantPandaCV

共 2792字,需浏览 6分钟

 ·

2022-02-13 14:10

前段时间逛GitHub看到FFCV这个库,该库主要是优化数据加载过程来提升整体训练速度。其中也放出了一些benchmark,看上去比其他优化库如DALI,PyTorch Lightening要快的不少。

一方面自己是搞框架的,数据加载优化是其中一部分重头戏;另一方面是PyTorch的数据加载速度也被诟病很久,毕竟面对的是研究人员,大部分人都是直接opencv, PIL一把梭哈数据预处理,我也很好奇如果好好写这部分能对PyTorch的速度提升多少,遂写这篇文章想分析分析(如有分析不对的地方还望指正)。

代码地址:https://github.com/libffcv/ffcv

使用文档:https://docs.ffcv.io/index.html

Reddit相关讨论:https://www.reddit.com/r/MachineLearning/comments/s781sr/p_ffcv_accelerated_model_training_via_fast_data/

快速上手

这里以提炼官方文档为主

制作数据集

https://docs.ffcv.io/writing_datasets.html

FFCV数据集是一个自定义格式.beton,所以第一步就需要将你的数据集转换成该格式。

这里我们以制作可索引数据集为例,首先创建一个支持索引的Dataset对象,你需要重写__getitem____len__方法

import numpy as np

class LinearRegressionDataset:
    def __init__(self, N, d):
        self.X = np.random.randn(N, d)
        self.Y = np.random.randn(N)

    def __getitem__(self, idx):
        return (self.X[idx].astype('float32'), self.Y[idx])

    def __len__(self):
        return len(self.X)

N, d = (1006)
dataset = LinearRegressionDataset(N, d)

这里创建了一个数据集,里面样本数量有100个,每个X维度为6,而Y维度为1

接着调用DatasetWriter将你的Dataset写成.beton格式

from ffcv.fields import NDArrayField, FloatField

writer = DatasetWriter(write_path, {
    'covariate': NDArrayField(shape=(d,), dtype=np.dtype('float32')),
    'label': FloatField(),

}, num_workers=16)
  • write_path 数据集要写入的路径
  • 字典,其中value项是你数据对应的一个Field对象。对应我们的数据集,每个X是一个ndarray,所以对应的是NDArrayField; 而Y则是一个浮点数,对应FloatField

使用DataLoader

制作好数据集我们就可以用了,这里的DataLoader其实是和PyTorch的很相似,使用方法如下

loader = Loader('/path/to/dataset.beton',
                batch_size=BATCH_SIZE,
                num_workers=NUM_WORKERS,
                order=ORDERING,
                pipelines=PIPELINES)
  • order 决定数据读取的顺序
  • pipelines 数据预处理的pipeline,我们可以把数据增广操作组合成一个pipeline传进来

pipeline一个组合示例如下:

image_pipeline: List[Operation] = [
    SimpleRGBImageDecoder(),
    RandomHorizontalFlip(),
    torchvision.transforms.ColorJitter(.4,.4,.4),
    RandomTranslate(padding=2),
    ToTensor(),
    ToDevice('cuda:0', non_blocking=True),
    ToTorchImage(),
    Convert(ch.float16),
    torchvision.transforms.Normalize(MEAN, STD), # Normalize using image statistics
])

至此简单介绍到这儿,我们来看下背后涉及到的一些技术

其构造主要分为以下几个大块:

- libffcv 自己写的一套C扩展
- ffcv python库主体
  |- fields 数据结构
  |- loader 数据加载器
  |- memory_manager 内存管理器
  |- pipeline 数据处理流水线
  |- transformer 增广操作
  |- traversal_order 数据遍历顺序控制

libffcv

作者基于Python C扩展写了一些必要的函数,包括如memcpy,fileread,imdecode,resize

其中resize使用的是Opencv来做,而图片解码采用的是turbojpeg库

fields

fields是ffcv里的数据结构,每个dataset的一个数据,是由一个或多个fields组成,每个field需要实现各自的编码,解码逻辑,分别对应数据集的写入,读取操作。

以FloatField为例:

class FloatField(Field):
    """
    A subclass of :class:`~ffcv.fields.Field` supporting (scalar) floating-point (float64)
    values.
    """

    def __init__(self):
        pass

    @property
    def metadata_type(self) -> np.dtype:
        return np.dtype(')

    @staticmethod
    def from_binary(binary: ARG_TYPE) -> Field:
        return FloatField()

    def to_binary(self) -> ARG_TYPE:
        return np.zeros(1, dtype=ARG_TYPE)[0]

    def encode(self, destination, field, malloc):
        destination[0] = field

    def get_decoder_class(self) -> Type[Operation]:
        return FloatDecoder

loader

ffcv loader对标 PyTorch DataLoader:

class Loader:
    def __init__(self,
                 fname: str,
                 batch_size: int,
                 num_workers: int = -1,
                 os_cache: bool = DEFAULT_OS_CACHE,
                 order: ORDER_TYPE = OrderOption.SEQUENTIAL,
                 distributed: bool = False,
                 seed: int = None,  # For ordering of samples
                 indices: Sequence[int] = None,  # For subset selection
                 pipelines: Mapping[str,
                                    Sequence[Union[Operation, ch.nn.Module]]] = {},
                 custom_fields: Mapping[str, Type[Field]] = {},
                 drop_last: bool = True,
                 batches_ahead: int = 3,
                 recompile: bool = False,  # Recompile at every epoch
                 )
:

我们挑几个重要的参数来说

  • os_cache 缓存策略
  • order 数据读取顺序
  • pipelines 数据预处理流水线,ffcv将所有的数据预处理集中到一个pipeline,然后借助JIT来加速相关处理操作
  • recompile 前面提到过他用JIT来加速预处理操作,当你每个epoch所对应的操作不一样,那么你就需要重新用JIT编译相关操作

memory_manager

这是一个内存管理对象,当数据集能够完全放进内存中时,则可以通过memory_manager设置相关策略,具体有两种策略。

一种是当内存充裕的时候,使用OS级别的cache,这里借助了np.memmap来完成虚拟内存和磁盘数据的映射,当出现缺页异常再执行相关的拷贝操作。

class OSCacheContext(MemoryContext):
    def __init__(self, manager:MemoryManager):
        self.manager = manager
        self.mmap = None

    @property
    def state(self):
        return (self.mmap, self.manager.ptrs, self.manager.sizes)

    def __enter__(self):
        res = super().__enter__()
        if self.mmap is None:
            self.mmap = np.memmap(self.manager.reader.file_name,
                                  'uint8', mode='r')
        return res

    # ...

另一种则是用进程级别的cache,维护固定数量的page,每一个batch释放相关的page,并对下一轮的数据进行预取prefetch。

# We now find how many pages we need to keep in our buffer     # We also determine where which page is going to reside    next_slot = 0    page_to_slot = {}    free_slots = set()    # For each batch    for b_id in range(len(pages_in_batch)):        # First we free the pages that are leaving        for page in leaving_at[b_id]:            free_slots.add(page_to_slot[page])        # We use the prefetch timing here because we want to be able        # To start prefetching ahead of time and not overwrite a slot        # That is currently used        for page in can_prefetch_at[b_id]:            # Then we find a slot for the incoming pages            if free_slots:                # There is a slot available for this page                slot = free_slots.pop()            else:                # We have to allocate a new slot because we ran out                slot = next_slot                next_slot += 1            page_to_slot[page] = slot    return Schedule(next_slot, page_to_slot,                    can_prefetch_at, entering_at, leaving_at)

Pipeline

里面具体有分了几个小部分

Operation

这是一个定义数据预处理操作的基类,其中generate_code方法用于返回相关处理操作的代码,以便后续被jit编译加速

class Operation(ABC):    def __init__(self):        self.matadata: np.ndarray = None        self.memory_read: Callable[[np.uint64], np.ndarray] = None        pass        # ...        @abstractmethod    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:         raise NotImplementedError

Compiler

顾名思义这是一个数据加载操作的"编译器",其思路就是利用numba.njit来将相关预处理操作编译,进行加速

class Compiler:    @classmethod    def set_enabled(cls, b):        cls.is_enabled = b    @classmethod    def set_num_threads(cls, n):        if n < 1 :            n = cpu_count()        cls.num_threads = n        set_num_threads(n)        ch.set_num_threads(n)    @classmethod    def compile(cls, code, signature=None):        parallel = False        if hasattr(code, 'is_parallel'):            parallel = code.is_parallel and cls.num_threads > 1                if cls.is_enabled:            return njit(signature, fastmath=True, nogil=True, error_model='numpy',                        parallel=parallel)(code)        return code

需要注意的是这里将fast_math默认开启,在一些浮点数的情形下可能会出现与普通计算不一致的情况(来自多年Loss对齐的惨痛教训)

然后我们看下 pipeline 主体代码,这是数据预处理的流水线,主要操作是:

  • 解析流水线

传进来的是一系列Operation的组合,需要先调用declare_state_and_memory来分配Operation对应的state和所需memory:

def parse_pipeline(self, batch_size=16):        memory_allocations: Mapping[int, Optional[Allocation]] = {}        operation_blocs = []        current_state: State = self.original_state        current_block = []        # We read the content of the pipeline, validate and collect        # Memory allocations        for op_id, operation in enumerate(self.operations):            previous_state = current_state            current_state, memory_allocation = operation.declare_state_and_memory(                current_state)            if current_state.jit_mode != previous_state.jit_mode:                if current_block:                    operation_blocs.append((previous_state.jit_mode, current_block))                current_block = [op_id]            else:                current_block.append(op_id)            memory_allocations[op_id] = memory_allocation        if current_block:            operation_blocs.append((current_state.jit_mode, current_block))        return operation_blocs, memory_allocations
  • 编译Operation代码

这部分很简单,就是逐个调用每个Operation的generate_code方法

def compile_ops(self):    compiled_ops = {}    for op_id, operation in enumerate(self.operations):        compiled_ops[op_id] = operation.generate_code()    return compiled_ops

这部分设计感觉是借鉴自NVIDIA DALI的Pipeline设计,FFCV这里借助了numba的jit特性,免去了大部分算子开发,只用JIT的特性就获取高性能,并且也易于用户在python端自定义拓展数据预处理操作。

Transform

这里是数据增广操作部分,通过继承Operation类,来重写generate_code逻辑。

以常用的ImageMixup为例:

class ImageMixup(Operation):    def __init__(self, alpha: float, same_lambda: bool):        super().__init__()        self.alpha = alpha        self.same_lambda = same_lambda    def generate_code(self) -> Callable:        alpha = self.alpha        same_lam = self.same_lambda        my_range = Compiler.get_iterator()        def mixer(images, dst, indices):            np.random.seed(indices[-1])            num_images = images.shape[0]            lam = np.random.beta(alpha, alpha) if same_lam else \                  np.random.beta(alpha, alpha, num_images)            for ix in my_range(num_images):                l = lam if same_lam else lam[ix]                dst[ix] = l * images[ix] + (1 - l) * images[ix - 1]            return dst        mixer.is_parallel = True        mixer.with_indices = True        return mixer    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:        return (previous_state, AllocationQuery(shape=previous_state.shape,                                                dtype=previous_state.dtype))

作者在Reddit上的一些讨论还提到了,他们实现了一个更快版本的NormalizeImage操作,对应的代码是在:https://github.com/libffcv/ffcv/blob/main/ffcv/transforms/normalize.py

实现具体分GPU和CPU版本,我们关注下GPU版本:

    def __init__(self, mean: np.ndarray, std: np.ndarray,                 type: np.dtype):        super().__init__()        table = (np.arange(256)[:, None] - mean[None, :]) / std[None, :]        # ...            def generate_code_gpu(self) -> Callable:        # We only import cupy if it's truly needed        import cupy as cp        import pytorch_pfn_extras as ppe        tn = np.zeros((), dtype=self.dtype).dtype.name        kernel = cp.ElementwiseKernel(f'uint8 input, raw {tn} table', f'{tn} output', 'output = table[input * 3 + i % 3];')        final_type = ch_dtype_from_numpy(self.original_dtype)        s = self        def normalize_convert(images, result):            B, C, H, W = images.shape            table = self.lookup_table.view(-1)            assert images.is_contiguous(memory_format=ch.channels_last), 'Images need to be in channel last'            result = result[:B]            result_c = result.view(-1)            images = images.permute(0, 2, 3, 1).view(-1)            current_stream = ch.cuda.current_stream()            with ppe.cuda.stream(current_stream):                kernel(images, table, result_c)            # Mark the result as channel last            final_result = result.reshape(B, H, W, C).permute(0, 3, 1, 2)            assert final_result.is_contiguous(memory_format=ch.channels_last), 'Images need to be in channel last'            return final_result.view(final_type)        return normalize_convert

这里的思路其实很巧妙,首先table是一个查找表,根据你传来的mean和std,提前计算了0-255这256个像素值经过归一化后的值。

比如 mean = [127.5, 127.5, 127.5], std = [1, 1, 1],那么得到的table shape为(256, 3),其中256代表着uint8像素值从0-255,而3代表的是RGB三个通道,数据为

[[-127.5 -127.5 -127.5# 像素值为0,RGB三个通道对应的normalized值 [-126.5 -126.5 -126.5] ...]

此时这个查找表是channel_last形式,我们用view把他展平:

table = self.lookup_table.view(-1)

基于表是channel_last形式,那对应的NCHW输入图片我们也要进行transpose,变成对应的NHWC并展平(我猜是为了后续访问连续,从而提升性能):

images = images.permute(0, 2, 3, 1).view(-1)

然后就可以调用cupy的ElementwiseKernel,进行逐元素操作:

kernel = cp.ElementwiseKernel(f'uint8 input, raw {tn} table'f'{tn} output''output = table[input * 3 + i % 3];')

其中input是输入像素值,i是index,这里对3取余得到具体是 RGB 3个通道中的哪一个。

总结

FFCV这个库还是挺不错的,不需要很多HPC知识,不需要你会写算子,通过比较成熟的一些工具来实现数据加载的加速,兼顾了PyTorch DataLoader的灵活性,同时又有较高的性能。

这个库到现在已经有1.5k star了,不得不说PyTorch的生态实在是好,基于其衍生出来的拓展库层出不穷。但也侧面反应出一些问题,需要依靠社区的力量来去完善。这个库给我们带来了很多新思路,有兴趣的朋友可以试试。

浏览 43
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报