首页 智能计算灵骏 用户指南 使用KSpeed优化套件 在PyTorch中使用KSpeed加速数据加载

在PyTorch中使用KSpeed加速数据加载

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,自定义方式步骤主要包括:

  1. 继承KSpeedDataset。

  2. 重定义__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操作
阿里云首页 智能计算灵骏 相关技术圈