Shortcuts

hfai.datasets

BaseDataset

hfai.dataset 基类

set_data_dir

设置数据集存放的主目录

get_data_dir

返回当前数据集主目录

download

下载数据集

ADE20k

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

CIFAR10

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

CIFAR100

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

CLUEForCLS

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

CLUEForMLM

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

CocoCaption

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

CocoDetection

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

CocoKeypoint

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

CocoPanoptic

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

ERA5

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

GoogleConceptualCaption

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

ImageNet

这是一个图像识别数据集

ImageNet21K

这是一个图像识别数据集

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.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。CIFAR10CIFAR100loader() 方法返回的是一个 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。CIFAR10CIFAR100loader() 方法返回的是一个 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个降水指标,按顺序依次是:u10v10t2mz@1000z@50z@500z@850mslr@500r@850spt@500t@850tcwvu@1000u@500u@850v@1000v@500v@850 和降水指标 tp

Parameters
  • split (str) – 数据集划分形式,包括:训练集(train)、验证集(val

  • mode (str) – 具体的训练模式,包括:pretrainfinetuneprecipitation

  • 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) – 数据集划分形式,包括:训练集(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.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) – 具体的数据名字,包括:ETTh1ETTh2ETTm1ETTm2exchange_rateelectricitynational_illnesstraffic

  • split (str) – 数据集划分形式,包括:训练集(train)或者验证集(val

  • seq_len (int) – 描述时间序列数据的序列长度

  • label_len (int) – 描述时间序列数据的标签长度

  • pred_len (int) – 描述时间序列数据的预测长度

  • features (str) – 预测的维度,包括:多维预测多维(M),单维预测单维(S),多维预测单维(MS

  • target (str) – 具体要预测的指标列,当 featuresS 或者 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": 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份
}

Return type

data_dict (dict)

Examples:

from hfai.datasets import NuScenes

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

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

for data_dict in loader:
    # training model
class hfai.datasets.OGB(data_name)[source]

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

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

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

  • 节点属性预测(Node Property Prediction): ogbn-productsogbn-proteinsogbn-arxivogbn-papers100Mogbn-mag

  • 连接预测(Link Property Prediction): ogbl-biokgogbl-citation2ogbl-collabogbl-ddiogbl-ddiogbl-ppaogbl-wikikg2

  • 图属性预测(Graph Property Prediction): ogbg-code2ogbg-molhivogbg-molpcbaogbg-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_temperaturerelative_humiditycomponent_of_windtotal_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