hfai.datasets¶

`BaseDataset`	hfai.dataset 基类
`set_data_dir`	设置数据集存放的主目录
`get_data_dir`	返回当前数据集主目录
`download`	下载数据集
`ADE20k`	这是一个语义分割的公开数据集
`AlphafoldData`	这是一个用于Alphafold训练的蛋白质预测数据集
`CIFAR10`	这是一个用于识别普适物体的小型数据集
`CIFAR100`	这是一个用于识别普适物体的大型数据集
`CLUEForCLS`	这是用于用于文本分类的 CLUE 数据集。
`CLUEForMLM`	这是一个用于掩蔽语言模型预训练的CLUE数据集
`CocoCaption`	这是一个用于图像说明的 COCO 数据集
`CocoDetection`	这是一个用于目标检测的 COCO 数据集
`CocoKeypoint`	这是一个用于关键点检测的COCO数据集
`CocoPanoptic`	这是一个用于全景分割的 COCO 数据集
`ERA5`	这是一个0.25°高分辨率的全球基础气象指标预报数据集
`GoogleConceptualCaption`	这是一个用于多模态训练的数据集
`ImageNet`	这是一个图像识别数据集
`ImageNet21K`	这是一个图像识别数据集
`ImageNet1K`	这是一个图像识别数据集
`ISC2021`	这是一个无监督学习的数据集
`KITTIObject2D`	这是一个目标检测数据集
`LJSpeech`	这是一个用于语音识别的数据集
`LTSF`	这是一个用于长时序预测的时间序列数据集
`NuScenes`	这是一个用于自动驾驶任务的 NuScenes 数据集
`OGB`	这是一个用于图机器学习的基准数据集
`WeatherBench`	这是一个用于天气预报的基准数据集

class hfai.datasets.BaseDataset[source]¶

hfai.dataset 基类

loader(*args, **kwargs)[source]¶

获取数据集的Dataloader

参数与 PyTorch的Dataloader 保持一致

Returns: 数据集的Dataloader

hfai.datasets.set_data_dir(path)[source]¶

设置数据集存放的主目录

我们会优先使用通过 set_data_dir 设置的路径，如果没有则会去使用环境变量 HFAI_DATASETS_DIR 的值。两者都没有设置的情况下，使用默认目录 /public_dataset/1/ffdataset。

Parameters: path (str, os.PathLike) – 数据集存放的主目录

Examples

>>> hfai.datasets.set_data_dir("/your/data/dir")
>>> hfai.datasets.get_data_dir()
PosixPath('/your/data/dir')

hfai.datasets.get_data_dir()[source]¶

返回当前数据集主目录

Returns: 当前数据集主目录
Return type: data_dir (Path)

Examples

>>> hfai.datasets.set_data_dir("/your/data/dir")
>>> hfai.datasets.get_data_dir()
PosixPath('/your/data/dir')

hfai.datasets.download(name, miniset=False, workers=4)[source]¶

下载数据集

Parameters

name (str) – 需要下载的数据集的名称，比如 CocoDetection
miniset (bool) – 是否使用 mini 集（默认为 False）
workers (int) – 进程数（默认为 4）

Returns

是否下载成功

Return type

success (bool)

Examples:

>>> hfai.datasets.set_data_dir('/your/data/dir')
>>> hfai.datasets.download("ImageNet", miniset=True)

class hfai.datasets.ADE20k(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个语义分割的公开数据集

该数据集拥有超过 25,000 张图像，这些图像用开放字典标签集密集注释。具体参考官网：https://groups.csail.mit.edu/vision/datasets/ADE20K/

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片和标注进行 transfrom，接受一张图片和一个 segmentation mask 作为输入，输出 transform 之后的图片和 segmentation mask
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个元组，包含一个RGB格式的图片，一个L格式的图片

Return type

pic, seg_mask (PIL.Image.Image, PIL.Image.Image)

Examples:

from hfai.datasets import ADE20k

def transform(pic, seg_mask):
    ...

dataset = ADE20k(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for pic, seg_mask in loader:
    # training model

class hfai.datasets.AlphafoldData(transform=None, check_data=True)[source]¶

这是一个用于Alphafold训练的蛋白质预测数据集

数据集中包含了131,291条蛋白质序列，一般整体作为训练集使用。其中pdb_mmcif和mmcif_cache_all.json数据在训练中不经dataset直接读取使用。

Parameters

transform (Callable) – transform函数，对蛋白质序列和alignment数据做处理
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）

Returns

返回一个元组，里面包含一个蛋白质的序列信息以及对应的alignment数据。

Return type

pdb_code, mmcif_string, bfd_hits, mgnify_hits, pdb70_hits, uniref90_hits

Examples:

from hfai.datasets import AlphafoldData
from torchvision import transforms

# 使用alphafold-optimized仓库时
transform = AlphafoldDataTransform()

dataset = AlphafoldData(transform)
loader = dataset.loader(batch_size=1, num_workers=8)

for data in loader:
    # training model

class hfai.datasets.CIFAR10(split, transform=None, target_transform=None)[source]¶

这是一个用于识别普适物体的小型数据集

该数据集一共包含 10 个类别的 RGB 彩色图片，每个图片的尺寸为 32 × 32 ，每个类别有 600 个图像，数据集中一共有 500 张训练图片和 100 张测试图片。更多信息参考官网：https://www.cs.toronto.edu/~kriz/cifar.html

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片进行 transfrom，接受一张图片作为输入，输出 transform 之后的图片
target_transform (Callable) – 对 target 进行 transfrom，接受一个 target 作为输入，输出 transform 之后的 target

Returns

返回的每条样本是一个元组，包含一个RGB格式的图片，及其对应的目标标签

Return type

image, target (PIL.Image.Image, int)

Examples:

from hfai.datasets import CIFAR10
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])

dataset = CIFAR10('train', transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for image, target in loader:
    # training model

Note

使用的时候所有数据会直接加载进内存，大小大约为 178 MiB。CIFAR10 和 CIFAR100 的 loader() 方法返回的是一个 torch.utils.data.DataLoader ，而不是 ffrecord.torch.DataLoader 。

class hfai.datasets.CIFAR100(split, transform=None, target_transform=None)[source]¶

这是一个用于识别普适物体的大型数据集

该数据集一共包含 100 个类别的 RGB 彩色图片，每个图片的尺寸为 32 × 32 ，每个类别有 600 个图像，数据集中一共有 500 张训练图片和 100 张测试图片。更多信息参考官网：https://www.cs.toronto.edu/~kriz/cifar.html

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片进行 transfrom，接受一张图片作为输入，输出 transform 之后的图片
target_transform (Callable) – 对 target 进行 transfrom，接受一个 target 作为输入，输出 transform 之后的 target

Returns

返回的每条样本是一个元组，包含一个 RGB 格式的图片，及其对应的目标标签

Return type

image, target (PIL.Image.Image, int)

Examples:

from hfai.datasets import CIFAR100
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])
dataset = CIFAR100('train', transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for image, target in loader:
    # training model

Note

使用的时候所有数据会直接加载进内存，大小大约为 178 MiB。CIFAR10 和 CIFAR100 的 loader() 方法返回的是一个 torch.utils.data.DataLoader ，而不是 ffrecord.torch.DataLoader 。

class hfai.datasets.CLUEForCLS(dataset_name, split)[source]¶

这是用于用于文本分类的 CLUE 数据集。

该数据集包含了多个子数据集：afqmc, cmnli, csl, iflytek, ocnli, tnews, wsc，可用来训练文本分类模型。

Parameters

dataset_name (str) – 子数据集的名字，比如 afqmc
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）

Returns

返回的每条样本是一个三元组，包括每个单词已经通过词汇表转成 id 的训练文本，单词掩蔽信息和对应的标签

Return type

text, pad_mask, label (torch.LongTensor, torch.BoolTensor, torch.LongTensor)

Examples:

from hfai.datasets import CLUEForMLM
dataset = CLUEForCLS("afqmc", split="train")
loader = dataset.loader(batch_size=32, num_workers=4)

for text, pad_mask, label in loader:
    # training model

class hfai.datasets.CLUEForMLM(check_data=True)[source]¶

这是一个用于掩蔽语言模型预训练的CLUE数据集

该数据集包含了三个 CLUE 子数据集，为训练 BERT 模型所定制，它们分别是：news, comment, web。该数据集共有 43926640 条语句。语句中有12%的单词会被 mask 掉，3% 的单词会随机替换成其他单词。

Parameters: check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
Returns: 返回的每条样本是一个三元组，包括每个单词已经通过词汇表转成 id 的训练文本，单词掩蔽信息和对应的标签
Return type: text, pad_mask, label (torch.LongTensor, torch.BoolTensor, torch.LongTensor)

Examples:

from hfai.datasets import CLUEForMLM
dataset = CLUEForMLM()
loader = dataset.loader(batch_size=32, num_workers=4)

for text, pad_mask, label in loader:
    # training model

class hfai.datasets.CocoCaption(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个用于图像说明的 COCO 数据集

更多信息参考：https://cocodataset.org

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片和标注进行 transfrom，接受一张图片、图片 id、标注图片和对应的标注作为输入，输出 transform 之后的图片、图片 id 和标注
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个三元组，包括一张 RGB 格式图片，对应的图片 ID，标注信息。标注如下例所示：

{
    'image_id': 444010,
    'id': 104057,
    'caption': 'A group of friends sitting down at a table sharing a meal.'
}

Return type

pic, id, anno (PIL.Image.Image, int, dict)

Examples:

from hfai.datasets import CocoCaption

def transform(pic, id, anno):
    ...

dataset = CocoCaption(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)
coco = dataset.coco # same as pycocotools.coco.COCO

for pic, id, anno in loader:
    # training model

class hfai.datasets.CocoDetection(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个用于目标检测的 COCO 数据集

更多信息参考：https://cocodataset.org

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片和标注进行 transfrom，接受一张图片、图片 id、标注图片和对应的标注作为输入，输出 transform 之后的图片、图片 id 和标注
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个三元组，包括一张 RGB 格式图片，对应的图片 ID，物体的标注信息。如下例所示：

{
    'segmentation': ...,
    'area': 3514.564,
    'iscrowd': 0,
    'image_id': 444010,
    'bbox': [x_left, y_top, w, h],
    'category_id': 44,
    'id': 91863
}

Return type

pic, id, anno (PIL.Image.Image, int, dict)

Examples:

from hfai.datasets import CocoDetection

def transform(pic, id, anno):
    ...

dataset = CocoDetection(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)
coco = dataset.coco # same as pycocotools.coco.COCO

for pic, id, anno in loader:
    # training model

class hfai.datasets.CocoKeypoint(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个用于关键点检测的COCO数据集

更多信息参考：https://cocodataset.org

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片和标注进行 transfrom，接受一张图片、图片 id、标注图片和对应的标注作为输入，输出 transform 之后的图片、图片 id 和标注
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个三元组，包括一张 RGB 格式图片，对应的图片 ID，物体的标注信息。标注如下例所示：

{
    'segmentation': ...,
    'num_keypoints': 11,
    'area': 34800.5498,
    'iscrowd': 0,
    'keypoints': ...,
    'image_id': 444010,
    'bbox': [x_left, y_top, w, h],
    'category_id': 1,
    'id': 1200757
}

Return type

pic, id, anno (PIL.Image.Image, int, dict)

Examples:

from hfai.datasets import CocoKeypoint

def transform(pic, id, anno):
    ...

dataset = CocoKeypoint(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)
coco = dataset.coco # same as pycocotools.coco.COCO

for pic, id, anno in loader:
    # training model

class hfai.datasets.CocoPanoptic(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个用于全景分割的 COCO 数据集

更多信息参考：https://cocodataset.org

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片和标注进行 transfrom，接受一张图片、图片 id、标注图片和对应的标注作为输入，输出 transform 之后的图片、图片 id 和标注
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个四元组，包括一张RGB格式图片，对应的图片ID，一张标注的RGB格式图片，物体的标注信息。如下例所示：

{
    "image_id": int,
    "file_name": str,
    "segments_info": {
        "id": int,
        "category_id": int,
        "area": int,
        "bbox": [x_left, y_top,width,height],
        "iscrowd": 0 or 1,
    },
}

Return type

pic, id, anno_pic, anno (PIL.Image.Image, int, PIL.Image.Image, dict)

Examples:

from hfai.datasets import CocoPanoptic

def transform(pic, id, anno_pic, anno):
    ...

dataset = CocoPanoptic(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)
coco = dataset.coco # same as pycocotools.coco.COCO

for pic, id, anno_pic, anno in loader:
    # training model

class hfai.datasets.ERA5(split, mode='pretrain', check_data=True, miniset=False)[source]¶

这是一个0.25°高分辨率的全球基础气象指标预报数据集

该数据集由欧洲中期天气预报中心（ECMWF）所构建并开源，从1979到2021年每6小时的全球气象数据。更多信息参考：https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5

数据包含20个气象指标和1个降水指标，按顺序依次是：u10， v10， t2m， z@1000， z@50， z@500， z@850， msl， r@500， r@850， sp， t@500， t@850， tcwv， u@1000， u@500， u@850， v@1000， v@500， v@850 和降水指标 tp

Parameters

split (str) – 数据集划分形式，包括：训练集（train）、验证集（val）
mode (str) – 具体的训练模式，包括：pretrain、finetune、precipitation
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每个样本是一个四元组，包括t时刻、t+1时刻和t+2时刻的气象指标数据，t+1时刻的降水数据

Return type

xt, xt1, xt2, pt1 (np.ndarray, np.ndarray, np.ndarray, np.ndarray)

Examples:

from hfai.datasets import ERA5

dataset = ERA5(split, mode='pretrain', check_data=True)
loader = dataset.loader(batch_size=3, num_workers=4)

for xt, xt1 in loader:
    # pretrain model

dataset = ERA5(split, mode='finetune', check_data=True)
loader = dataset.loader(batch_size=1, num_workers=4)

for xt, xt1, xt2 in loader:
    # finetune model

dataset = ERA5(split, mode='precipitation', check_data=True)
loader = dataset.loader(batch_size=4, num_workers=4)

for xt, pt1 in loader:
    # training precipitation model

get_scaler()[source]¶

获取WeatherBench数据的统计特征信息。

Returns: 数据分布统计对象，指标顺序如主类描述，包含：指标数据的均值（mean），指标数据的方差（std）

class hfai.datasets.GoogleConceptualCaption(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个用于多模态训练的数据集

该数据集是一个子数据集，从 3318333 个 “图片-字幕” 对中随机采样了 2850879 个。更多信息参考：https://ai.google.com/research/ConceptualCaptions/

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片和文本进行 transfrom，接受一张图片和一段文本作为输入，输出 transform 之后的结果
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每个样本是一个元组，包含一个RGB格式的图片，一段字幕文本

Return type

pic, text (PIL.Image.Image, str)

Examples:

from hfai.datasets import GoogleConceptualCaption
from torchvision import transforms

img_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])
tokenize = ...

def transform(pic, text):
    pic = img_transform(pic)
    text = tokenize(text)
    return pic, text

dataset = GoogleConceptualCaption(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for pic, text in loader:
    # training model

class hfai.datasets.ImageNet(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个图像识别数据集

更多信息参考：https://image-net.org

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片进行 transfrom，接受一张图片作为输入，输出 transform 之后的图片
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每个样本是一个元组，包括一个RGB格式的图片，以及代表这张图片类别的标签

Return type

pic, label (PIL.Image.Image, int)

Examples:

from hfai.datasets import ImageNet
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])
dataset = ImageNet(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for pic, label in loader:
    # training model

class hfai.datasets.ImageNet21K(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个图像识别数据集

更多信息参考：https://github.com/Alibaba-MIIL/ImageNet21K

Parameters

split (str) – 数据集划分形式，包括：完整集（full）、训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片进行 transfrom，接受一张图片作为输入，输出 transform 之后的图片
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每个样本是一个元组，包括一个RGB格式的图片，图片名，以及代表这张图片类别的标签

Return type

pic, name, label (PIL.Image.Image, str, int)

Examples:

from hfai.datasets import ImageNet21K
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])
dataset = ImageNet21K(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for pic, name, label in loader:
    # training model

class hfai.datasets.ImageNet1K(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个图像识别数据集

更多信息参考：https://github.com/Alibaba-MIIL/ImageNet21K

Parameters

split (str) – 数据集划分形式，包括：完整集（full）、训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对图片进行 transfrom，接受一张图片作为输入，输出 transform 之后的图片
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每个样本是一个元组，包括一个RGB格式的图片，图片名，以及代表这张图片类别的标签

Return type

pic, name, label (PIL.Image.Image, str, int)

Examples:

from hfai.datasets import ImageNet1K
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])
dataset = ImageNet1K(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for pic, name, label in loader:
    # training model

class hfai.datasets.ISC2021(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个无监督学习的数据集

该数据集是 Facebook AI 在 NeurIPS 2021 举办的图像相似检索大赛上所开源的百万量级的数据集。更多信息参考：https://github.com/facebookresearch/isc2021

Parameters

split (str) – 数据集划分形式，包括：训练集（trainining）、参考集（reference）或者查询集（query）
transform (Callable) – transform 函数，对图片进行 transfrom，接受一张图片作为输入，输出 transform 之后的图片
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一张RGB格式的图片

Return type

img (PIL.Image.Image)

Examples:

from hfai.datasets import ISC2021
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])
dataset = ISC2021(split, transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for imgs in loader:
    # training model

class hfai.datasets.KITTIObject2D(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个目标检测数据集

KITTI 数据集由德国卡尔斯鲁厄理工学院和丰田美国技术研究院联合创办，是目前国际上最大的自动驾驶场景下的计算机视觉算法评测数据集。更多信息参考：http://www.cvlibs.net/datasets/kitti/

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者测试集（test）
transform (Callable) – transform 函数，对图片和标注进行 transfrom，接受一张图片和一个 target 作为输入，输出 transform 之后的图片和 target。（测试集没有 target）
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个元组，包含一个RGB格式的图片，以及对应的包含位置信息的字典，例如：

{
    'type': 'Pedestrian',
    'truncated': 0.0,
    'occluded': 0,
    'alpha': -0.2,
    'bbox': [x1, y1, x2, y2],
    'dimensions': [1.89, 0.48, 1.2],
    'location': [1.84, 1.47, 8.41],
    'rotation_y': 0.01
}

Return type

pic, target (PIL.Image.Image, list)

Examples:

from hfai.datasets import KITTIObject2D

def transform(pic, target):
    ...

dataset = KITTIObject2D('train', transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for pic, target in loader:
    # training model

class hfai.datasets.LJSpeech(transform=None, check_data=True, miniset=False)[source]¶

这是一个用于语音识别的数据集

该数据集由 13100 个来自 7 部非小说类书籍的单个说话者阅读段落的短音频片段组成。更多信息参考：https://keithito.com/LJ-Speech-Dataset

Parameters

transform (Callable) – transform 函数，对样本进行 transfrom，接受一条样本作为输入，输出 transform 之后的样本
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个五元组，包括原始 wav 文件的 id，采样率，wav 文件的数据，音频对应的文本，归一化后的文本

Return type

wav_id, sample_rate, wavefrom, text, norm_text (str, int, np.ndarray, str, str)

Examples:

from hfai.datasets import LJSpeech

def transform(wav_id, sample_rate, wavefrom, text, norm_text):
    ...

dataset = LJSpeech(transform)
loader = dataset.loader(batch_size=64, num_workers=4)

for wav_id, sample_rate, wavefrom, text, norm_text in loader:
    # training model

class hfai.datasets.LTSF(data_name, split, seq_len=96, label_len=48, pred_len=24, features='S', target='OT', timeenc=0, freq='h')[source]¶

这是一个用于长时序预测的时间序列数据集

该数据集包含电力、经济、医疗、交通等不同领域下的长期时间序列数据。更多信息参考：https://github.com/thuml/Autoformer

Parameters

data_name (str) – 具体的数据名字，包括：ETTh1，ETTh2，ETTm1，ETTm2，exchange_rate，electricity，national_illness，traffic
split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
seq_len (int) – 描述时间序列数据的序列长度
label_len (int) – 描述时间序列数据的标签长度
pred_len (int) – 描述时间序列数据的预测长度
features (str) – 预测的维度，包括：多维预测多维（M），单维预测单维（S），多维预测单维（MS）
target (str) – 具体要预测的指标列，当 features 为 S 或者 MS 时奏效
timeenc (int) – 不同的时间特征编码，默认为 0
freq (str) – 时间特征编码的频率, 包括：每秒（s），每分钟（t），每小时（h），每天（d），每工作日（b），每星期（w），每月（m）。你也可以选择更定制的频次，例如：15min 或者 3h

Returns

返回的每个样本是一个四元组，包括历史指标序列，未来指标序列，历史时间位置信息编码，未来时间位置信息编码

Return type

seq_x, seq_y, seq_x_mask, seq_y_mask (np.ndarray, np.ndarray, np.ndarray, np.ndarray)

Examples:

from hfai.datasets import LTSF

dataset = LTSF(data_name, split, seq_len, label_len, pred_len, features, target, timeenc, freq)
loader = dataset.loader(batch_size=64, num_workers=4)

for seq_x, seq_y, seq_x_mask, seq_y_mask in loader:
    # training model

get_scaler()[source]¶

获取LTSF数据的统计特征信息。

Returns: 数据分布统计对象，包含：指标数据的均值（mean），指标数据的方差（std）

class hfai.datasets.NuScenes(split, transform=None, check_data=True, miniset=False)[source]¶

这是一个用于自动驾驶任务的 NuScenes 数据集

更多信息参考：https://www.nuscenes.org/

Parameters

split (str) – 数据集划分形式，包括：训练集（train）或者验证集（val）
transform (Callable) – transform 函数，对数据字典进行 transfrom，接受数据字典作为输入，输出 transform 之后的数据字典
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）
miniset (bool) – 是否使用 mini 集合（默认为 False）

Returns

返回的每条样本是一个数据字典，如下例所示 (参考官网数据解释)：

{
    "imgs": list of PIL.Image.Image, # 六个相机的视觉图像
    "trans": torch.Tensor, # 机身坐标系下，六个相机的偏移向量
    "rots": torch.Tensor, # 机身坐标系下，六个相机的旋转向量
    "intrins": torch.Tensor, # 六个相机的内参
    "lidar_data": torch.Tensor, # 机身坐标系下，激光雷达数据
    "lidar_mask": torch.Tensor, # 标识激光雷达填充数据的掩码
    "car_trans": torch.Tensor, # 世界坐标系下，机身偏移向量
    "yaw_pitch_roll": torch.Tensor, # 世界坐标系下，机身的旋转欧拉角
    "semantic_gt": torch.Tensor, # 语义分割标注
    "instance_gt": torch.Tensor, # 实例分割标注
    "direction_gt": torch.Tensor # 方向分割标注，将360度划分为36份
    "sweeps": list of dict # 非没有标注的中间帧的传感器数据
    "gt_boxes": torch.Tensor # 7 个自由度的 3D 包围框，一个 Nx7 数组
    "gt_names": list of string # 3D 包围框的类别，一个 1xN 数组
    "gt_velocity": torch.Tensor # 3D 包围框的速度（由于不准确，没有垂直测量），一个 Nx2 数组
    "num_lidar_pts": torch.Tensor # 每个 3D 包围框中包含的激光雷达点数
    "num_radar_pts": torch.Tensor # 每个 3D 包围框中包含的雷达点数
    "valid_flag": torch.Tensor # 每个包围框是否有效。将包含至少一个激光雷达或雷达点的 3D 框作为有效框
}

Return type

data_dict (dict)

Note

detection 相关的 key，每个样本的数据长度不同。例如 ‘gt_boxes’，第 0 条数据长度为 10，而第 1 条数据长度为 20。

Examples:

from hfai.datasets import NuScenes
from torchvision import transforms
import torch

def transform(data_dict):
    trans = transforms.Compose([transforms.ToTensor()])
    data_dict['imgs'] = torch.stack([
        trans(img) for img in data_dict['imgs']
    ])
    return data_dict

split = 'train'  # or val
dataset = NuScenes(split, transform)

# segmentation_task
loader = dataset.loader(batch_size=64, num_workers=4)
for i, data_dict in enumerate(loader):
    print('keys', data_dict.keys())
'''
output: keys dict_keys([
'imgs', 'trans', 'rots', 'intrins', 'lidar_data', 'lidar_mask', 'car_trans',
'yaw_pitch_roll', 'semantic_gt', 'instance_gt', 'direction_gt'
])
'''

# detection_task
dataset.set_task('detection')
for i, data_dict in enumerate(loader):
    print('detection keys', data_dict.keys())

'''
output: keys dict_keys([
'imgs', 'trans', 'rots', 'intrins',
'lidar_data', 'lidar_mask', 'car_trans', 'yaw_pitch_roll', 'semantic_gt',
'instance_gt', 'direction_gt', 'sweeps', 'gt_boxes', 'gt_names',
 'gt_velocity', 'num_lidar_pts', 'num_radar_pts', 'valid_flag'
 ])
'''

class hfai.datasets.OGB(data_name)[source]¶

这是一个用于图机器学习的基准数据集

Open Graph Benchmark (OGB) 是一组真实的、大规模的、多样化的基准数据集，用于在图上进行机器学习。更多信息参考：https://ogb.stanford.edu/

该数据集收纳了如下图数据：

节点属性预测（Node Property Prediction）： ogbn-products，ogbn-proteins，ogbn-arxiv，ogbn-papers100M，ogbn-mag
连接预测（Link Property Prediction）： ogbl-biokg，ogbl-citation2，ogbl-collab，ogbl-ddi，ogbl-ddi，ogbl-ppa，ogbl-wikikg2
图属性预测（Graph Property Prediction）： ogbg-code2，ogbg-molhiv，ogbg-molpcba，ogbg-ppa
大规模图（Large-Scale Chanllenge）： ogblsc-wikikg2

Parameters: data_name (str) – 数据的名字
Returns: 返回的是 torch_geometric 格式的图数据，和“训练、验证、测试”样本分割索引
Return type: data, split (torch_geometric.data.data.Data, dict)

Examples

加载 ogbn- 系列数据集

>>> from hfai.datasets import OGB
>>> dataset = OGB(data_name='ogbn-proteins')
>>> dataset.get_data()
Data(num_nodes=132534, edge_index=[2, 79122504], edge_attr=[79122504, 8], node_species=[132534, 1], y=[132534, 112])

加载 ogbl- 系列数据集

>>> from hfai.datasets import OGB
>>> dataset = OGB(data_name='ogbl-ppa')
>>> dataset.get_data()
Data(num_nodes=576289, edge_index=[2, 42463862], x=[576289, 58])

加载 ogbg- 系列数据集

>>> from hfai.datasets import OGB
>>> dataset = OGB(data_name='ogbg-molhiv')
>>> dataset.get_data()
[Data(edge_index=[2, 40], edge_attr=[40, 3], x=[19, 9], y=[1, 1], num_nodes=19),
 Data(edge_index=[2, 88], edge_attr=[88, 3], x=[39, 9], y=[1, 1], num_nodes=39),
 ...

加载 ogblsc- 系列数据集

>>> from hfai.datasets import OGB
>>> dataset = OGB(data_name='ogblsc-wikikg2')
>>> dataset.get_data()
{'num_entities': 91230610, 'num_relations': 1387, 'num_feat_dims': 768, 'entity_feat': ..., 'realtion_feat': ..., 'train_hrt': ..., 'val_hr': ..., 'val_t': ...}

get_data()[source]¶

获取OGB图数据。

Returns: 节点/链接预测任务下返回torch_geometric格式的图数据，图属性预测任务下是返回元组，大规模图任务下返回具有ndarray格式数据的字典

get_split()[source]¶

获取OGB图数据的分割索引。

Returns: 返回字典，包含“训练、验证、测试”样本分割的索引。大规模图任务下该方法不可用

loader(*args, **kwargs)[source]¶: OGB数据不支持直接Dataloader。

class hfai.datasets.WeatherBench(data_name, split, include_context=False, check_data=True)[source]¶

这是一个用于天气预报的基准数据集

该数据集基于欧洲气象局标准数据集ERA5采样所得，从1979到2019年每小时的全球气象数据。更多信息参考：https://github.com/pangeo-data/WeatherBench

Parameters

data_name (str) – 具体的气象预报项，包括：2m_temperature、relative_humidity、component_of_wind、total_cloud_cover
split (str) – 数据集划分形式，包括：训练集（train）、验证集（val）或者测试集（test）
include_context (bool) – 是否包含气象上下文信息作为输入
check_data (bool) – 是否对每一条样本检验校验和（默认为 True）

Returns

返回的每个样本是一个二元组，包括历史指标序列数据，未来指标序列

Return type

seq_x, seq_y (np.ndarray, np.ndarray)

Examples:

from hfai.datasets import WeatherBench

dataset = WeatherBench(data_name, split, include_context, check_data)
loader = dataset.loader(batch_size=64, num_workers=4)

for seq_x, seq_y in loader:
    # training model

get_kernel()[source]¶

获取WeatherBench数据的地理核信息。

Returns: 核信息是一个四元组，包含：稀疏矩阵（sparse_idx），局部卷积核输入（MLP_inputs），测地线（geodesic），切面角度比（angle_ratio）
Return type: kernel (tuple)

get_scaler()[source]¶

获取WeatherBench数据的统计特征信息。

Returns: 数据分布统计对象，包含：指标数据的均值（mean），指标数据的方差（std）