数据管理平台 SDK 操作指南

版本记录

版本名称	功能描述	发布时间
数据管理平台 GA 版 v1.1.0	SDK 操作指南从原始长文档中拆分出来，使文档具备更好的可读性。	2023-4-28
数据管理平台 GA 版 v1.3.0	支持通过SDK加载本地数据集同名Python脚本，实现远程同名数据集的访问管理。	2023-6-30
数据管理平台 GA 版 v1.5.0	支持通过SDK加载指定格式(tar/json/csv/parquet)文件列表	2023-8-30
数据管理平台 GA 版 v1.6.0	支持数据空间	2023-9-30

SDK加载数据集

load_dataset

加载远端 repo 的指定数据集

SDK 工具可以进行一键式安装，支持 python3.7 及以上版本，下面为详细安装的命令：

pip install https://aids-download.aoss.cn-sh-01.sensecoreapi-oss.cn/packages/pypi/gravdataset.tar.gz

安装了 SDK 工具 gravdataset 后，导入 load_dataset 的模块 ，使用 load_dataset 函数，就可以连接远端的 repo 数据集，使用 PyTorch 的 dataloader 进行训练。

from gravdataset import load_dataset
from torch.utils.data import DataLoader
dataset = load_dataset(<data_space_name/repo_name>, access_key_id=<YOUR_ACCESS_KEY_ID>, access_key_secret=<YOUR_ACCESS_KEY_SECRET>)
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

py脚本加载

load_dataset 函数默认调用指定 repo 名的同名 python 文件，所加载的 py 脚本主要包含三部分，一定义数据集的信息结构，二划分不同的数据子集，比如训练集、测试集、验证集，三生成相应的样本。

load_dataset 加载数据时，不会将数据集的原始数据下载到本地存储，而是直接去数据集所在的对象存储拿相应的原始数据。

默认会加载 master 分支的最新版本，同时支持 revision 参数，revision 可以指定相应 repo 特定的分支，比如 master、dev 等，指定分支时为相应分支的最新版本；也可以指定相应 repo 特定的版本 commit ID ，7 位以上均可以支持，即为相对应的 commit 版本。

示例如下：

from gravdataset import load_dataset
from torch.utils.data import DataLoader
dataset = load_dataset(<data_space_name/repo_name>, revision=dev, access_key_id=<YOUR_ACCESS_KEY_ID>, access_key_secret=<YOUR_ACCESS_KEY_SECRET>)
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

支持 split 参数，split 可以指定相应的划分数据集，如 train、test、validation，如果给定，则返回相对应的单个数据集，比如只返回给定的 train 数据集，如果不给定，则以字典的形式返回所有的数据集，使用时比如为 DataSet.split.train 。

支持 name 参数，一些数据集的 py 脚本中，划定了各种数据子集，比如每个子数据集包含不同语言类型的数据。这些子数据集必须在加载数据集时明确选择一个，如果不提供配置名称 name ，则报错提醒选择一个配置名称 name 。

支持 loadScript 参数，可以指定加载脚本的本地地址（优先本地加载脚本），通过本地的加载脚本，加载远端的数据

from gravdataset import load_dataset 
load_dataset("OpenDataSet/ImageNet1K", access_key_id="XXX", access_key_secret="XXX", load_script="./ImageNet1K.py")

如果相对应的 AKSK 错误，请前往AccessKey访问秘钥页面处核实。

对于数据管理平台提供的公开数据集，也可以采用 load_dataset 方式加载数据集，用于训练，示例如下：

from gravdataset import load_dataset
from torch.utils.data import DataLoader
dataset = load_dataset("OpenDataSet/ImageNet1K",access_key_id="XXX", access_key_secret="XXX")
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

指定文件加载

load_dataset方法也支持通过data_dir和data_files参数加载指定tar/json/parquet/csv格式的文件，示例如下

from gravdataset import load_dataset
from torch.utils.data import DataLoader
# 同时加载val/train的文件列表
data_files = {"val":["val_0.csv","val_1.csv"],'train':["train_0.csv","train_1.csv"]}
# # 加载文件列表
# data_files = ["train_0.csv","train_1.csv"]
# # 加载单独文件
# data_files = "train_0.csv"
# # 加载一批文件，会加载train_0.csv train_0.csv ... train_xxx.csv文件 
# data_files = "train_*.csv"
data_dir = "en"
dataset = load_dataset("OpenDataSet/ImageNet1K",access_key_id="XXX", access_key_secret="XXX", data_files=data_files,data_dir=data_dir)
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

加载CC3M tar文件示例

from gravdataset import load_dataset
from torch.utils.data import DataLoader
# 同时加载val/train的文件列表
data_files = "*.tar"
data_dir = "data/00000"
dataset = load_dataset("OpenDataSet/CC3M",access_key_id="XXX", access_key_secret="XXX", data_files=data_files,data_dir=data_dir)
for i in dataset:
    print(i)     # {"__tarname__":"xx.tar","__key__":"123214",json:"<json-content>","image":"<image_bytes>","txt":"this is picture"}

py脚本撰写说明

通过脚本加载数据集

我们可以通过编写 python 脚本的方式自定义数据集加载方式，包括数据集的基本元信息、数据集的划分和配置、数据集处理下载和生成数据集。该脚本位于与数据集相同的文件夹或存储库中，并且应与数据集名称一致，如下所示：

my_dataset/
├── README.md
├── my_dataset.py
└── data/

以该结构存储的数据集可以使用以下命令来加载：

from gravdataset import load_dataset
dataset = load_dataset("path/to/my_dataset")

数据集加载脚本编写教程

本教程将展示如何为图像数据集创建数据集加载脚本，主要包含以下内容：

• 数据集脚本的基本结构
• 定义数据集元信息
• 定义数据集划分规则
• 生成数据集

1. 数据集脚本的基本结构

一个数据集脚本必须包含一个具体的 DatasetBuilder 类（一般继承于常用的 GeneratorBasedBuilder），该类定义了我们生成数据集所需的四部分操作：数据集配置、数据集元信息、数据集划分规则、数据集生成器，具体模板如下：

import gravdataset

class COCO(gravdataset.GeneratorBasedBuilder):
    """COCO dataset."""

    BUILDER_CONFIGS = [
        gravdataset.BuilderConfig(
            name='COCO2014',
            version=VERSION,
            description='COCO2014 dataset for det and segm'),
        gravdataset.BuilderConfig(
            name='COCO2017',
            version=VERSION,
            description='COCO2017 dataset for det and segm'),
    ]

    DEFAULT_CONFIG_NAME = 'COCO2017'

    def _info(self):

    def _split_generators(self, dl_manager):

    def _generate_examples(self, images, metadata_path):

以上类模板包含两个属性变量和三个函数接口需要我们自定义实现，具体含义如下:

• BUILDER_CONFIGS（可缺省）：数据集配置变量，是一个 list 类型，这里我们一个数据集脚本是可以有多个数据集在内的，该列表内的每一个配置代表一个数据集，主要配置数据集的名称（name）、版本（version）、描述（description）等，比如以上模板表示了 COCO 的 2014 和 2017 两个版本的数据集。用户可通过 load_dataset 函数的 name 参数来指定具体的数据集配置

• DEFAULT_CONFIG_NAME（可缺省）：默认的数据集配置，用于选择如果加载数据集时未提供 name 参数，默认使用的数据集配置

• _info：该函数定义了数据集的元信息，该信息在生成的 dataset 内，可以随时访问

• _split_generators：该函数内执行数据集下载和数据集划分，一般用于划分训练、测试、验证集

• _generate_examples：该函数定义一个数据生成器，每次迭代返回一条具体的数据

2. 定义数据集元信息(_info)

添加数据集的元信息有助于用户了解数据集的具体信息。该信息存储在由 _info 方法返回的 DatasetInfo 类中。用户可以通过以下方式访问此信息：

from gravdataset import load_dataset
ds = load_dataset("/path/to/COCO", name='COCO2017', split='train', access_key_id=ak, access_key_secret=sk)
print(ds.info)

你可以指定很多数据集信息，一般常用以下字段：

• description：提供数据集的简明描述
• features：指定数据字段类型
• supervised_keys：指定输入的字段
• home：提供指向数据集主页的链接
• citation：是数据集的 BibTeX 引用
• license：数据集的许可证
• meta_info：额外的元数据信息，一般用于指定数据集的类别信息，如 meta_info=dict(classes=['cat', 'dog'])

其中最常用的信息字段为 description、features、meta_info，如 COCO 数据集的 _info 函数示例如下：

import gravdataset
from gravdataset.features import Features, Value, Sequence

def _info(self):
    return gravdataset.DatasetInfo(
        description='COCO dataset for detection and instance segmentation task.',
        # _CLASSES 指 COCO 数据集的 80 个类别
        meta_info=dict(classes=_CLASSES),
        # feature 可缺省，若未传入则根据 _generate_examples 函数生成的数据结构自动生成
        features=Features({
            'img_info': {
                'filename': Value('string'),
                'height': Value('int32'),
                'width': Value('int32'),
            },
            'ann_info': {
                'bboxes': Sequence(Sequence(Value('float64'))),
                'labels': Sequence(Value('int64')),
                'masks': Sequence(Sequence(Sequence(Value('float64')))),
                'bboxes_ignore': Sequence(Sequence(Value('float64'))),
                'label_ignore': Sequence(Value('int64')),
                'seg_map': Value('string')
            }
        }))

3. 下载并定义数据集划分规则(_split_generators)

_split_generators 函数一般包含两部分内容，数据标注文件下载和数据集划分。

• dl_manager：该参数为下载器，用于下载数据标注文件，该下载器包含两个方法 dl_manager.download() 和 dl_manager.download_and_extract() 分别用于下载文件和下载并解压文件。

• SplitGenerator：下载标注文件后，使用 SplitGenerator 指定划分数据集的 name 和相应的数据集参数 gen_kwargs，gen_kwargs 是一个 dict 类型变量，其中的 key 会作为参数传入 _generate_examples 函数。

具体的实现示例如下：

# 这里的 _URLS 指定了数据集中的图片和标注文件的相对路径
_URLS = {
    'COCO2014': {
        'train_prefix': 'train2014',
        'train_meta': 'annotations/instances_train2014.json',
        'val_prefix': 'val2014',
        'val_meta': 'annotations/instances_val2014.json'
    },
    'COCO2017': {
        'train_prefix': 'train2017',
        'train_meta': 'annotations/instances_train2017.json',
        'val_prefix': 'val2017',
        'val_meta': 'annotations/instances_val2017.json'
    },
}

def _split_generators(self, dl_manager):
    # 这里的 self.config 为之前 BUILDER_CONFIGS 中的数据集配置，根据 load_dataset 传入的 name 参数指定
    train_prefix = _URLS[self.config.name]['train_prefix']
    train_meta = _URLS[self.config.name]['train_meta']
    val_prefix = _URLS[self.config.name]['val_prefix']
    val_meta = _URLS[self.config.name]['val_meta']
    # 下载标注文件
    train_meta = dl_manager.download(train_meta)
    val_meta = dl_manager.download(val_meta)
    return [
        gravdataset.SplitGenerator(
            name='train',
            # These kwargs will be passed to _generate_examples
            gen_kwargs={
                'img_prefix': train_prefix,
                'ann_file': train_meta
            }),
        gravdataset.SplitGenerator(
            name='val',
            # These kwargs will be passed to _generate_examples
            gen_kwargs={
                'img_prefix': val_prefix,
                'ann_file': val_meta
            }),
    ]

4. 生成数据集(_generate_examples)

该函数指定生成数据的具体格式，输入参数为 gen_kwargs 的 key 字段。

该函数必须是一个生成器，返回数据编号 index 和具体的返回数据 sample。

一般在该函数内通过标注文件或文件列表构造模型所需的具体输入数据，如 COCO 数据集可生成检测或实例分割所需的数据，具体示例如下：

from pycocotools.coco import COCO

def _generate_examples(self, img_prefix, ann_file):
    """Parser coco format annotation file."""
    coco = COCO(ann_file)
    cat_ids = coco.getCatIds(_CLASSES)
    cat2label = {cat_id: i for i, cat_id in enumerate(cat_ids)}
    img_ids = coco.getImgIds()
    index = 0
    for i in img_ids:
        sample = dict(img_info=dict())
        info = coco.loadImgs([i])[0]
        sample['img_info']['filename'] = os.path.join(
            img_prefix, info['file_name'])
        sample['img_info']['height'] = info['height']
        sample['img_info']['width'] = info['width']
        ann_ids = coco.getAnnIds([i])
        ann_info = coco.loadAnns(ann_ids)
        gt_bboxes = []
        gt_labels = []
        gt_bboxes_ignore = []
        gt_label_ignore = []
        gt_masks_ann = []
        for i, ann in enumerate(ann_info):
            if ann.get('ignore', False):
                continue
            x1, y1, w, h = ann['bbox']
            inter_w = max(0, min(x1 + w, info['width']) - max(x1, 0))
            inter_h = max(0, min(y1 + h, info['height']) - max(y1, 0))
            if inter_w * inter_h == 0:
                continue
            if ann['area'] <= 0 or w < 1 or h < 1:
                continue
            if ann['category_id'] not in cat_ids:
                continue
            bbox = [x1, y1, x1 + w, y1 + h]
            if ann.get('iscrowd', False):
                gt_bboxes_ignore.append(bbox)
                gt_label_ignore.append(cat2label[ann['category_id']])
            else:
                gt_bboxes.append(bbox)
                gt_labels.append(cat2label[ann['category_id']])
                gt_masks_ann.append(ann.get('segmentation', None))

        seg_map = sample['img_info']['filename'].rsplit('.', 1)[0] + '.png'

        sample['ann_info'] = dict(
            bboxes=gt_bboxes,
            labels=gt_labels,
            bboxes_ignore=gt_bboxes_ignore,
            label_ignore=gt_label_ignore,
            masks=gt_masks_ann,
            seg_map=seg_map)
        yield index, sample
        index += 1

至此我们就将一个数据集脚本所需的要素全部实现，只需将该脚本和数据上传至 SenseCore 数据管理平台，即可在线 load_dataset。

如何使用脚本加载数据集

通过示例我们可以看到返回的数据仅包含数据集的标注数据和原始图片的路径，我们这里具体迭代数据集时需要通过 Dataset 的 get_file 函数读取具体的图片。

具体示例如下：

import io
from PIL import Image
from gravdataset import load_dataset

ds = load_dataset(
    'aidmpdev/COCO',
    name='COCO2017',
    split='val',
    access_key_id=ak,
    access_key_secret=sk)

print(f'categories info: {ds.info.meta_info["classes"]}')

data_info = ds[0]
img_path = data_info['img_info']['file_name']
img_buffer = ds.get_file(img_path)
img = Image.open(io.BytesIO(img_buffer))
img.show()

脚本示例

下面为公开数据集 COCO 的同名 python 脚本撰写示例：

import os

from pycocotools.coco import COCO
import gravdataset
from gravdataset.features import Features, Value, Sequence

_DESCRIPTION = 'COCO dataset for detection and instance segmentation task.'

_URLS = {
    'COCO2014': {
        'train_prefix': 'train2014',
        'train_meta': 'annotations/instances_train2014.json',
        'val_prefix': 'val2014',
        'val_meta': 'annotations/instances_val2014.json'
    },
    'COCO2017': {
        'train_prefix': 'train2017',
        'train_meta': 'annotations/instances_train2017.json',
        'val_prefix': 'val2017',
        'val_meta': 'annotations/instances_val2017.json'
    },
}

_CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
            'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
            'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
            'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
            'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
            'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
            'scissors', 'teddy bear', 'hair drier', 'toothbrush')


class Coco(gravdataset.GeneratorBasedBuilder):
    """COCO dataset for detection and instance segmentation task."""

    VERSION = gravdataset.Version('0.1.0')

    BUILDER_CONFIGS = [
        gravdataset.BuilderConfig(
            name='COCO2014',
            version=VERSION,
            description='COCO2014 dataset for det and segm'),
        gravdataset.BuilderConfig(
            name='COCO2017',
            version=VERSION,
            description='COCO2017 dataset for det and segm'),
    ]
    # It's not mandatory to have a default configuration.
    # Just use one if it make sense.
    DEFAULT_CONFIG_NAME = 'train'

    def _info(self):
        return gravdataset.DatasetInfo(
            # This is the description that will appear on the datasets page.
            description=_DESCRIPTION + f'\nCLASSES: ({",".join(_CLASSES)})',
            features=Features({
                'img_info': {
                    'filename': Value('string'),
                    'height': Value('int32'),
                    'width': Value('int32'),
                },
                'ann_info': {
                    'bboxes': Sequence(Sequence(Value('float64'))),
                    'labels': Sequence(Value('int64')),
                    'masks': Sequence(Sequence(Sequence(Value('float64')))),
                    'bboxes_ignore': Sequence(Sequence(Value('float64'))),
                    'label_ignore': Sequence(Value('int64')),
                    'seg_map': Value('string')
                }
            }))

    def _split_generators(self, dl_manager):
        train_prefix = _URLS[self.config.name]['train_prefix']
        train_meta = _URLS[self.config.name]['train_meta']
        val_prefix = _URLS[self.config.name]['val_prefix']
        val_meta = _URLS[self.config.name]['val_meta']
        train_meta = dl_manager.download(train_meta)
        val_meta = dl_manager.download(val_meta)
        return [
            gravdataset.SplitGenerator(
                name='train',
                # These kwargs will be passed to _generate_examples
                gen_kwargs={
                    'img_prefix': train_prefix,
                    'ann_file': train_meta
                }),
            gravdataset.SplitGenerator(
                name='val',
                # These kwargs will be passed to _generate_examples
                gen_kwargs={
                    'img_prefix': val_prefix,
                    'ann_file': val_meta
                }),
        ]

    def _generate_examples(self, img_prefix, ann_file):
        """Parser coco format annotation file."""
        coco = COCO(ann_file)
        cat_ids = coco.getCatIds(_CLASSES)
        cat2label = {cat_id: i for i, cat_id in enumerate(cat_ids)}
        img_ids = coco.getImgIds()
        index = 0
        for i in img_ids:
            sample = dict(img_info=dict())
            info = coco.loadImgs([i])[0]
            sample['img_info']['filename'] = os.path.join(
                img_prefix, info['file_name'])
            sample['img_info']['height'] = info['height']
            sample['img_info']['width'] = info['width']
            ann_ids = coco.getAnnIds([i])
            ann_info = coco.loadAnns(ann_ids)
            gt_bboxes = []
            gt_labels = []
            gt_bboxes_ignore = []
            gt_label_ignore = []
            gt_masks_ann = []
            for i, ann in enumerate(ann_info):
                if ann.get('ignore', False):
                    continue
                x1, y1, w, h = ann['bbox']
                inter_w = max(0, min(x1 + w, info['width']) - max(x1, 0))
                inter_h = max(0, min(y1 + h, info['height']) - max(y1, 0))
                if inter_w * inter_h == 0:
                    continue
                if ann['area'] <= 0 or w < 1 or h < 1:
                    continue
                if ann['category_id'] not in cat_ids:
                    continue
                bbox = [x1, y1, x1 + w, y1 + h]
                if ann.get('iscrowd', False):
                    gt_bboxes_ignore.append(bbox)
                    gt_label_ignore.append(cat2label[ann['category_id']])
                else:
                    gt_bboxes.append(bbox)
                    gt_labels.append(cat2label[ann['category_id']])
                    gt_masks_ann.append(ann.get('segmentation', None))

            seg_map = sample['img_info']['filename'].rsplit('.', 1)[0] + '.png'

            sample['ann_info'] = dict(
                bboxes=gt_bboxes,
                labels=gt_labels,
                bboxes_ignore=gt_bboxes_ignore,
                label_ignore=gt_label_ignore,
                masks=gt_masks_ann,
                seg_map=seg_map)
            yield index, sample
            index += 1