如何使用KSpeed客户端以及提供的数据加载模块进行Pytorch计算框架中的数据加载加速-阿里云帮助中心

KSpeed库除了提供客户端和kspeedcmd命令行工具之外，还提供了一个工具集utils，包含了基于Pytorch定制优化的数据加载模块。本文为您介绍如何使用KSpeed客户端以及提供的数据加载模块进行Pytorch计算框架中的数据加载加速。

KSpeed数据加载加速方案概览

KSpeed通过数据预加载和数据预取两级Pipeline实现数据加载加速。KSpeed首先将数据集预加载到分布式的计算集群各节点上，然后利用计算侧Zero-Copy的RDMA高性能数据传输进行数据预取，并提供Cast-Free的Python API，加速AI训练中对共享数据I/O的加载性能。用户可以根据需求和实际情况使用以下两种模式使用KSpeed加速数据加载。

基于主动预加载的数据加载加速

基于主动预加载的数据加载加速，由训练脚本在数据集Dataset初始化时自动预加载全量数据集到分布式的计算集群节点，在运行时调用KSpeed API进行数据访问预取，在训练结束时，数据集自动从KSpeed空间卸载。此种加载方式KSpeed主动实现数据集的预加载和卸载。

基于离线预加载的数据加载加速

基于离线预加载的数据加载加速，由用户使用命令行工具kspeedcmd离线地将全量数据集预加载到分布式计算集群KSpeed空间中，在运行时调用KSpeed API进行数据访问预取。基于离线预加载的数据加载加速，在训练结束时，数据集不会自动从KSpeed空间卸载，训练任务反复共享该离线加载的数据集。

Pytorch数据加载加速

Pytorch的数据加载包括三个重要的类：Dataset、Dataloader、Sampler，典型的Pytorch数据加载框架如下：

#定义dataset
train_ds = Dataset()
#定义dataloader（使用默认的sampler）
dataloader = torch.utils.data.DataLoader(train_ds, batch_size=args.batch_size, num_workers=args.dlworker)

在Pytorch计算框架中使用KSpeed，最少仅需要改写Dataset，将Dataset实现为基于KSpeed的高性能数据加载。KSpeed提供了一个数据工具包kspeed.utils，其中提供了KSpeedDataset类模板和几种典型的Dataset模块，如KSpeedDataset、KSpeedImageFolder。实现基于KSpeed数据读取的Dataset途径有：

调用kspeed.utils里面提供的现成的Dataset，如KSpeedDataset、KSpeedImageFolder等。
用户继承kspeed.utils里面提供的基类KSpeedDataset构建自定义的Dataset。

说明

KSpeedDataset类模板实现了数据集自动预加载功能，在KSpeedDataset类的初始化时会进行数据集是否已经预加载到kspeed服务集群进行检查判断：

如果数据集已经预加载到kspeed服务集群，则不进行数据主动预加载，训练结束后也不删除kspeed服务集群中的数据集。
如果数据集未预加载到kspeed服务集群，则进行数据主动预加载，训练结束后自动从kspeed服务集群中删除数据集。

使用KSpeedImageFolder模块

说明

此种方式主要适用于CV深度学习图片数据集加载加速场景。

kspeed.utils包含的KSpeedImageFolder是面向CV场景定制的，用户可以直接调用KSpeedImageFolder类，传入图片数据集的存储根路径以及图片预处理transform组合，然后结合Pytorch的DataLoader和Sampler实现高性能的数据加载模型。使用方法如下：

#import KSpeedImageFolder
from kspeed.utils.data.kspeeddataset import KSpeedImageFolder

#定义dataset
train_ds = KSpeedImageFolder(root=root, transform=None, target_transform=None)
#定义sampler、dataloader
if torch.distributed.is_initialized():
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_ds)
else:
    train_sampler = None
dataloader = torch.utils.data.DataLoader(train_ds, sampler = train_sampler, num_workers=dlworker)

实例化KSpeedImageFolder类对象的参数描述如下：

参数	说明
root	指定文件集合数据集的根路径
objname	指定文件集合数据集在KSpeed空间的对象名字, 默认None，如果不提供则默认是使用root的顶层目录名字
dl_workers	指定dataloader中配置使用加载数据的worker数目。因为kspeed是基于RDMA的数据加载服务，RDMA在父子多进程使用时有特别限制，该输入参数要指定KSpeedImageFolder是运行在主进程还是子进程，直接传入dataloader的workers数目即可： dl_workers=0，则表明KSpeedImageFolder是运行在主进程 dl_workers>0，否则表明KSpeedImageFolder是运行在子进程，默认dl_worker=4，表明是运行在子进程
kspeed_iplist	指定了连接kspeed server所在的iplist, 在K8s环境可以自动发现KSpeed服务，可以不用输入，默认None
username	指定了连接kspeed server的用户名, 默认admin
password	指定了连接kspeed server的用户名密码, 默认admin
transform	指定图片预处理操作，默认None
target_transform	指定图片target的预处理操作，默认None

使用KSpeedDataset模块

说明

此种方式主要适用于NLP深度学习结构化数据集加载加速场景。

kspeed.utils包含的KSpeedDataset是面向Array数据集场景定制的，使用KSpeedDataset构建pytorch dataloader的方式如下：

#import KSpeedDataset
from kspeed.utils.data.kspeeddataset import KSpeedDataset

#定义dataset
train_ds = KSpeedDataset(kspeed_iplist, filename, batch_size, in_main_process)
#定义sampler、dataloader
batch_sampler_tr = BatchSampler(RandomSampler(train_ds),batch_size=batch_size, drop_last=True)
dataloader = torch.utils.data.DataLoader(train_ds, sampler = batch_sampler_tr, batch_size=None, num_workers=dlworker)

实例化KSpeedDataset类对象的参数描述如下：

参数	说明
root	指定文件集合数据集的根路径
objname	指定文件集合数据集在KSpeed空间的对象名字, 默认None，如果不提供则默认是使用root的顶层目录名字
dl_workers	指定dataloader中配置使用加载数据的worker数目。因为kspeed是基于RDMA的数据加载服务，RDMA在父子多进程使用时有特别限制，该输入参数要指定KSpeedImageFolder是运行在主进程还是子进程，直接传入dataloader的workers数目即可： dl_workers=0，则表明KSpeedImageFolder是运行在主进程 dl_workers>0，否则表明KSpeedImageFolder是运行在子进程，默认dl_worker=4，表明是运行在子进程
device	指定数据加载CPU还是GPU显存，默认CPU
batch_size	指定了KSpeedDataset读取的数据集的batch_size 。KSpeed提供两种获取数据的方式：非batch模式（输入batch_size=1），给定一个index返回一个sample的numpy batch模式（输入batch_size=batch_size），给定一个batch_size的index list，返回一个batch samples的numpy
kspeed_iplist	指定了连接kspeed server所在的iplist, 在K8s环境可以自动发现KSpeed服务，可以不用输入，默认None
username	指定了连接kspeed server的用户名, 默认admin
password	指定了连接kspeed server的用户密码, 默认admin

自定义Dataset模块使用KSpeed

说明

此种方式可根据使用场景提供客户自定义服务。

用户还可以继承KSpeedDataset自定义Pytorch Dataset，自定义方式步骤主要包括：

继承KSpeedDataset。
重定义__getitem__(self, index)方法。

from kspeed.kspeedcluster import kspeedclient
from kspeed.utils.data.kspeeddataset import KSpeedDataset

class CostomKSpeedDataset(KSpeedDataset):
    def __init__(self,
                 root,
                 object_name=None,
                 dl_workers=4,
                 kspeed_iplist=None,
                 username='admin',
                 password='admin'):
        super(CostomKSpeedDataset,
              self).__init__(root, object_name, dl_workers, 'CPU', 1,
                             kspeed_iplist, username, password)

    def __getitem__(self, index):
        if self.dataset is None:
            self.load()
        #重定义getitem操作