安装所有额外功能

pip install 'litdata[extras]'

混合搭配不同的数据集以进行实验并创建更好的模型。

使用CombinedStreamingDataset合并数据集。例如，这种Slimpajama和StarCoder的混合在TinyLLAMA项目中被用来预训练一个1.1B的Llama模型，训练了3万亿个词元。

from litdata import StreamingDataset, CombinedStreamingDataset, StreamingDataLoader, TokensLoader
from tqdm import tqdm
import os

train_datasets = [
    StreamingDataset(
        input_dir="s3://tinyllama-template/slimpajama/train/",
        item_loader=TokensLoader(block_size=2048 + 1), # LLM使用的优化词元加载器
        shuffle=True,
        drop_last=True,
    ),
    StreamingDataset(
        input_dir="s3://tinyllama-template/starcoder/",
        item_loader=TokensLoader(block_size=2048 + 1), # LLM使用的优化词元加载器
        shuffle=True,
        drop_last=True,
    ),
]

# 按以下比例混合SlimPajama数据和Starcoder数据：
weights = (0.693584, 0.306416)
combined_dataset = CombinedStreamingDataset(datasets=train_datasets, seed=42, weights=weights, iterate_over_all=False)

train_dataloader = StreamingDataLoader(combined_dataset, batch_size=8, pin_memory=True, num_workers=os.cpu_count())

# 遍历合并后的数据集
for batch in tqdm(train_dataloader):
    pass

使用train_test_split将数据集分割为训练、验证、测试集。

from litdata import StreamingDataset, train_test_split

dataset = StreamingDataset("s3://my-bucket/my-data") # 数据存储在云端

print(len(dataset)) # 显示数据的长度
# 输出: 100,000

train_dataset, val_dataset, test_dataset = train_test_split(dataset, splits=[0.3, 0.2, 0.5])

print(train_dataset)
# 输出: 30,000

print(val_dataset)
# 输出: 20,000

print(test_dataset)
# 输出: 50,000

  处理数据的一个较小、可管理的部分，以节省时间和资源。

from litdata import StreamingDataset, train_test_split

dataset = StreamingDataset("s3://my-bucket/my-data", subsample=0.01) # 数据存储在云端

print(len(dataset)) # 显示数据的长度
# 输出: 1000

向现有数据集添加新数据或在需要时重新开始，提供灵活的数据管理。

LitData优化后的数据集被假定为不可变的。但是，你可以通过将模式更改为append或overwrite来决定修改它们。

from litdata import optimize, StreamingDataset

def compress(index):
    return index, index**2

if __name__ == "__main__":
    # 添加一些数据
    optimize(
        fn=compress,
        inputs=list(range(100)),
        output_dir="./my_optimized_dataset",
        chunk_bytes="64MB",
    )

    # 稍后，你添加更多数据
    optimize(
        fn=compress,
        inputs=list(range(100, 200)),
        output_dir="./my_optimized_dataset",
        chunk_bytes="64MB",
        mode="append",
    )

    ds = StreamingDataset("./my_optimized_dataset")
    assert len(ds) == 200
    assert ds[:] == [(i, i**2) for i in range(200)]

overwrite模式将删除现有数据并重新开始。

通过使用高级压缩算法减少数据占用空间。

import litdata as ld

def compress(index):
    return index, index**2

if __name__ == "__main__":
    # 添加一些数据
    ld.optimize(
        fn=compress,
        inputs=list(range(100)),
        output_dir="./my_optimized_dataset",
        chunk_bytes="64MB",
        num_workers=1,
        compression="zstd"
    )

使用zstd，你可以在这个简单示例中达到4.34倍的高压缩率。

无需下载整个数据集或将其加载到本地机器上，即可查看大型数据集的特定部分。

from litdata import StreamingDataset

dataset = StreamingDataset("s3://my-bucket/my-data") # 数据存储在云端

print(len(dataset)) # 显示数据的长度

print(dataset[42]) # 显示数据集的第42个元素

自定义数据处理方式以更好地满足你的需求。

继承StreamingDataset并重写其__getitem__方法来添加任何额外的数据转换。

from litdata import StreamingDataset, StreamingDataLoader
import torchvision.transforms.v2.functional as F

class ImagenetStreamingDataset(StreamingDataset):

    def __getitem__(self, index):
        image = super().__getitem__(index)
        return F.resize(image, (224, 224))

dataset = ImagenetStreamingDataset(...)
dataloader = StreamingDataLoader(dataset, batch_size=4)

for batch in dataloader:
    print(batch.shape)
    # 输出: (4, 3, 224, 224)

测量并优化数据加载速度，提高效率。

StreamingDataLoader支持对数据加载过程进行分析。只需使用profile_batches参数指定要分析的批次数：

from litdata import StreamingDataset, StreamingDataLoader

StreamingDataLoader(..., profile_batches=5)

这将生成一个名为result.json的Chrome跟踪文件。然后，通过在Chrome浏览器中打开chrome://tracingURL并加载该跟踪文件来可视化这个跟踪。

高效处理大型数据文件,不会过度占用计算机内存。

处理大型文件(如压缩的 parquet 文件)时,使用 Python 的 yield 关键字一次处理和存储一个项目,从而减少整个程序的内存占用。

from pathlib import Path
import pyarrow.parquet as pq
from litdata import optimize
from tokenizer import Tokenizer
from functools import partial

# 1. 定义一个函数,将 parquet 文件中的文本转换为标记
def tokenize_fn(filepath, tokenizer=None):
    parquet_file = pq.ParquetFile(filepath)
    # 按批次处理以减少内存使用
    for batch in parquet_file.iter_batches(batch_size=8192, columns=["content"]):
        for text in batch.to_pandas()["content"]:
            yield tokenizer.encode(text, bos=False, eos=True)

# 2. 生成输入
input_dir = "/teamspace/s3_connections/tinyllama-template"
inputs = [str(file) for file in Path(f"{input_dir}/starcoderdata").rglob("*.parquet")]

# 3. 将优化后的数据存储在 "/teamspace/datasets" 或 "/teamspace/s3_connections" 下的任何位置
outputs = optimize(
    fn=partial(tokenize_fn, tokenizer=Tokenizer(f"{input_dir}/checkpoints/Llama-2-7b-hf")), # 注意:使用 HF tokenizer 或其他任何 tokenizer
    inputs=inputs,
    output_dir="/teamspace/datasets/starcoderdata",
    chunk_size=(2049 * 8012), # 每个块存储的标记数。这大约是每个块 64MB 的标记。
)

限制临时文件使用的磁盘空间,防止存储问题。

调整 StreamingDataset 的本地缓存限制。这有助于确保下载的数据块在使用后被删除,并保持磁盘使用量较低。

from litdata import StreamingDataset

dataset = StreamingDataset(..., max_cache_size="10GB")

指定缓存文件应存储的目录,确保高效的数据检索和管理。这对于组织数据存储和提高访问时间特别有用。

from litdata import StreamingDataset
from litdata.streaming.cache import Dir

cache_dir = "/path/to/your/cache"
data_dir = "s3://my-bucket/my_optimized_dataset"

dataset = StreamingDataset(input_dir=Dir(path=cache_dir, url=data_dir))

优化本地网络上计算机的数据处理,以提高现场设置的性能。

本地计算节点可以挂载和使用网络驱动器。网络驱动器是本地局域网上的共享存储设备。为了减少网络负载,StreamingDataset 支持对数据块进行 缓存。

from litdata import StreamingDataset

dataset = StreamingDataset(input_dir="local:/data/shared-drive/some-data")

Lightning 可以将大型工作负载并行分配到数百台机器上。通过扩展到足够多的机器,可以将完成数据处理任务的时间从数周缩短到几分钟。

要在多台机器上应用优化操作,只需向其提供 num_nodes 和 machine 参数,如下所示:

import os
from litdata import optimize, Machine

def compress(index):
    return (index, index ** 2)

optimize(
    fn=compress,
    inputs=list(range(100)),
    num_workers=2,
    output_dir="my_output",
    chunk_bytes="64MB",
    num_nodes=2,
    machine=Machine.DATA_PREP, # 你可以在几十种优化的机器之间选择
)

如果 output_dir 是本地路径,优化后的数据集将位于: /teamspace/jobs/{job_name}/nodes-0/my_output。否则,它将存储在指定的 output_dir 中。

读取优化后的数据集:

from litdata import StreamingDataset

output_dir = "/teamspace/jobs/litdata-optimize-2024-07-08/nodes.0/my_output"

dataset = StreamingDataset(output_dir)

print(dataset[:])

同时对数据集的不同部分应用相同的更改,以节省时间和精力。

map 操作符可用于在输入列表上应用函数。

这里是一个使用 map 操作符在一个大图像文件夹上应用 resize_image 函数的示例。

from litdata import map
from PIL import Image

# 注意:输入也可以直接引用 s3 上的文件。
input_dir = "my_large_images"
inputs = [os.path.join(input_dir, f) for f in os.listdir(input_dir)]

# resize_image 函数接收一个输入(image_path)和输出目录。
# 写入 output_dir 的文件将被持久化。
def resize_image(image_path, output_dir):
  output_image_path = os.path.join(output_dir, os.path.basename(image_path))
  Image.open(image_path).resize((224, 224)).save(output_image_path)

map(
    fn=resize_image,
    inputs=inputs,
    output_dir="s3://my-bucket/my_resized_images",
)

使用不同的云存储服务,提供数据存储灵活性和成本节省选项。

将 S3 兼容的对象存储服务器(如 MinIO)与 litdata 集成,非常适合本地基础设施设置。使用环境变量或配置文件配置端点和凭证。

设置环境变量以连接到 MinIO:

export AWS_ACCESS_KEY_ID=access_key
export AWS_SECRET_ACCESS_KEY=secret_key
export AWS_ENDPOINT_URL=http://localhost:9000  # MinIO 端点

或者,在 ~/.aws/{credentials,config} 中配置凭证和端点:

mkdir -p ~/.aws && \
cat <<EOL >> ~/.aws/credentials
[default]
aws_access_key_id = access_key
aws_secret_access_key = secret_key
EOL

cat <<EOL >> ~/.aws/config
[default]
endpoint_url = http://localhost:9000  # MinIO 端点
EOL

在 LitData with MinIO 仓库中探索 litdata 与 MinIO 的示例设置,以了解实际实现细节。

通过对单个样本或块应用加密来保护您的数据,确保敏感信息在存储过程中受到保护。

这个示例演示了如何使用 FernetEncryption 类进行样本级加密的数据优化函数。

from litdata import optimize
from litdata.utilities.encryption import FernetEncryption
import numpy as np
from PIL import Image
# 使用密码初始化FernetEncryption进行样本级加密
fernet = FernetEncryption(password="your_secure_password", level="sample")
data_dir = "s3://my-bucket/optimized_data"

def random_image(index):
    """生成一个随机图像用于演示目的。"""
    fake_img = Image.fromarray(np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8))
    return {"image": fake_img, "class": index}

# 优化数据并应用加密
optimize(
    fn=random_image,
    inputs=list(range(5)),  # 示例输入：[0, 1, 2, 3, 4]
    num_workers=1,
    output_dir=data_dir,
    chunk_bytes="64MB",
    encryption=fernet,
)

# 将加密密钥保存到文件以供日后使用
fernet.save("fernet.pem")

您可以使用StreamingDataset类加载加密数据，如下所示：

from litdata import StreamingDataset
from litdata.utilities.encryption import FernetEncryption

# 加载加密密钥
fernet = FernetEncryption(password="your_secure_password", level="sample")
fernet.load("fernet.pem")

# 创建一个流式数据集用于读取加密样本
ds = StreamingDataset(input_dir=data_dir, encryption=fernet)

如果您想实现自己的加密方法，可以继承Encryption类并定义必要的方法：

from litdata.utilities.encryption import Encryption

class CustomEncryption(Encryption):
    def encrypt(self, data):
        # 在此实现您的自定义加密逻辑
        return data

    def decrypt(self, data):
        # 在此实现您的自定义解密逻辑
        return data

通过这种设置，您可以确保数据安全的同时保持加密处理的灵活性。

• Imagenet-1.2M数据集包含1,281,167张图像。
• 为了与其他基准测试保持一致，我们测量了从AWS S3加载的几个框架的流式传输速度（每秒图像数）。