关于深度学习:基于PySpark的10亿级数据集LAION5B元数据快速处理实践全文分享

举荐语

多模态大模型训练热火朝天开展，但以 LAION-5B 为代表的大规模多模态数据集获取却成了一个令人头疼的问题。

OpenDataLab 两位工程师在浦数 AI Talk 做了十分实用的 LAION-5B 下载教训分享，咱们整顿了其演讲内容、Parquet 文件、图片下载工具，心愿能对大家下载同类数据集提供帮忙和参考。以下为全文内容：

一、数据集背景

依据历史钻研发现，随着训练数据减少时，ALIGN、BASIC、Turing Bletchly、FLORENCE 和 GLIDE 等大型多模态视觉语言模型在新的短少样本标签的数据集上也有很强的迁徙能力，而且性能还在稳固进步。但这些模型须要数十亿的图文数据才有达到良好的成果，到 2022 年为止，还没有十亿规模的公开图文对数据集。

直到 LAION-5B 公布，该数据集由 5.85Billoin CLIP 过滤的图像文本对组成，它为多模态预训练提供十分重要的“燃料”。（之前咱们写过 LAION-5B 数据集解读，戳此回顾：https://zhuanlan.zhihu.com/p/571741834）

二、LAION-5B 数据集构造

依据官网文件统计，LAION-5B 数据有 5,860,068,373 个样本，依照语言被官网划分为 3 个子数据集，别离是：

1. laion2b-en 2.32 billion of these contain texts in the English language
2. laion2b-multi 2.26 billion contain texts from 100+ other languages
3. laion1b-nolang 1.27 billion have texts where a particular language couldn’t be clearly detected.

其中每个数据集官网提供了原始图片的 URL，能够依据 URL 下载图片文件，以及些 URL 上的标签。
这部分元数据被存储在 parquet 文件中。样例 parquet 文件构造如下：

data_sample
├── laion2B-en
│   ├── part-00006-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
│   ├── part-00014-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
│   ├── part-00039-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
│   ├── part-00043-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
│   ├── part-00078-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
│   ├── part-00093-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
│   └── part-00123-5114fd87-297e-42b0-9d11-50f1df323dfa-c000.snappy.parquet
└── laion2B-multi
    ├── part-00001-fc82da14-99c9-4ff6-ab6a-ac853ac82819-c000.snappy.parquet
    ├── part-00026-fc82da14-99c9-4ff6-ab6a-ac853ac82819-c000.snappy.parquet
    ├── part-00030-fc82da14-99c9-4ff6-ab6a-ac853ac82819-c000.snappy.parquet
    ├── part-00034-fc82da14-99c9-4ff6-ab6a-ac853ac82819-c000.snappy.parquet
    └── part-00125-fc82da14-99c9-4ff6-ab6a-ac853ac82819-c000.snappy.parquet

三、parquet 元数据处理

在官网下载 parquet 元数据时，发现以下几个小问题：

1. similarity、aesthetic_score 等指标散布在多个 parquet 文件中，字段扩散、类型不对立，须要屡次下载。应用时须要先关联组合查问，TB 级的文件处理速度慢，须要高配置的服务器进行解决；
2. parquet 文件中图片存储门路规定不明确，通过 parquet 过滤筛选图片时，无奈关联下载图片的存储门路和其它字段
3. parquet 文件中 parquet_id、hash 等字段反复，影响图片的惟一索引
4. 通过 url 下载的图片格式未知（有 webp、jpg、png、avif 等多种格局），影响下载图片的预览和存储

为了满足不同场景的数据应用需要，保障图片惟一索引 ID，咱们对官网的 parquet 文件进行了关联合并、字段补充等操作，造成一张字段丰盛的“宽表”，数据表构造与字段设计如下：

上表中，最初一列是 parquet 文件起源，示意字段对应的 parquet 文件。
这里应用了官网的 3 处 parquet 文件，数据预览、下载链接如下：

1. 初始 Laion5B
   https://huggingface.co/datasets/laion/laion2B-en
   https://huggingface.co/datasets/laion/laion2B-multi
   https://huggingface.co/datasets/laion/laion1B-nolang
2. Joined: with punsafe and pwatermark
   https://huggingface.co/datasets/laion/laion2B-en-joined
   https://huggingface.co/datasets/laion/laion2B-multi-joined
   https://huggingface.co/datasets/laion/laion1B-nolang-joined
3. Laion-aesthetic
   Laion aesthetic is a laion5B subset with aesthetic > 7 pwatermark < 0.8 punsafe < 0.5 See
   https://huggingface.co/datasets/laion/laion1B-nolang-aesthetic
   https://huggingface.co/datasets/laion/laion2B-en-aesthetic
   https://huggingface.co/datasets/laion/laion2B-multi-aesthetic

四、解决流程及步骤

上面聊聊“宽表”的加工处理过程，有需要的共事可参考对官网的原始 parquet 进行解决。嫌麻烦的同学，能够交给 opendatalab，在网站下载解决好的 parquet 文件。https://opendatalab.com/LAION-5B

因为 parquet 文件数据量较大，有几个 TB，这里咱们应用了大数据集群进行了分布式解决。
● 应用的技术栈有：
Spark/Hadooop/Hive/HDFS/Impala
● 集群硬件配置：
服务器 3 台，48core Cpu, 750GB Memory, 4TB Hard disk
● 数据处理过程和流程图如下：
数据输出：

• 下载官网 parquet 文件，并 load 到 Hive 表
• 解析下载的图片，判断图片类型，造成 id, image_path, image_suffix 的映射文件

数据处理：

• 读取 Hive 表数据，通过 PySpark 对 Hive 表的数据进行分布式 join 关联操作
  数据输入：
• Hive 后果表导出为 parquet 格式文件，并上传至 OSS/Ceph 存储

为了不便数据处理，这里对数据表进行简略的分层：
– ODS 层：
原始 parquet 文件 load Hive 后的结构化数据表，其中表 2 是对表 1 字段进行了裁减，表 3 是下载图片相干的信息。因为官网 parquet 文件只提供了下载 url 链接，咱们并不知道图片类型和后缀，所以对下载的图片文件进行程序断定，辨认出图片类型，对应 image_suffix 字段，image_path 是图片的存储门路。
– DMD 层：
通过对表 2、3、4 进行 join 关联操作，生成两头表 6
– DMS 层：
将两头表 6 与含有 punsafe、pwatermark 信息的表 5 进行关联，失去最初的后果表 7

数据处理操作和代码示例如下：

4-1. Data load

次要操作是将 parquet 文件 load 到 Hive 表，load 操作实现后，失去图中的 1、3、4、5 四张 Hive 表。
以初始 parquet 文件 load 为例，示例代码如下。

import os
from pyspark.sql import HiveContext, SQLContext
from pyspark.sql.functions import lit, input_file_name
from pyspark.sql.functions import col, udf
from pyspark.sql.types import StringType, LongType
import mmh3

sc = spark.sparkContext
sql_context = SQLContext(sc)

# 通过 url/text 计算 hash 值 
def compute_hash(url, text):
    if url is None:
        url = ''
        
    if text is None:
        text = ''total = (url + text).encode("utf-8")
    return mmh3.hash64(total)[0]

# 注册 spark udf
udf_compute_hash = udf(compute_hash, LongType())

# 提取 input_file_name 门路中的文件名称
def path_proc(file_path):
    return str(file_path).split("/")[-1]

udf_path_proc = udf(path_proc)

# 因数据总量较大，这里按子集分批读取
parquet_path = "/nvme/datasets/laion5b/parquet/laion2B-en"
# parquet_path = "/nvme/datasets/laion5b/parquet/laion2B-multi"
# parquet_path = "/nvme/datasets/laion5b/parquet/laion2B-nolang"

# Hive 一级分区名称
head_tail = os.path.split(parquet_path)
partition_name = head_tail[1]

parquet_df =  spark.read.parquet(f"file://{parquet_path}/*.parquet")
parquet_df = parquet_df.withColumn("file_name", input_file_name())
parquet_df = parquet_df.withColumn("hash", udf_compute_hash(parquet_df["URL"], parquet_df["TEXT"]))
parquet_df = parquet_df.withColumn("dir_name",  lit(f"{partition_name}"))
parquet_df = parquet_df.withColumn("file_name", udf_path_proc(parquet_df["file_name"]))

# 自定义视图名称，并注册视图
view_name = 'parquet_view'
parquet_df.createOrReplaceTempView(f"{view_name}")

# 数据写入 Hive 分区表，一级分区名称 dir_name，二级分区名称 file_name
sql_context.sql(f"insert overwrite table laion5b.parquet_view partition(dir_name='{partition_name}', file_name) select SAMPLE_ID,URL,TEXT,HEIGHT,WIDTH, LICENSE,NSFW,similarity,`hash`, file_name from {view_name}")

4-2. Data processing

数据处理过程次要包含数据表裁减、hash join 操作。

因为表 1 的数据量较大，存在字段冗余，这里对表 1 的局部字段进行裁减失去表 2。
表 2、3、4 的 join 代码如下，先将图片的 sample_id、licenese、nsfw、image_suffix、aesthetic_score 字段，按 hash 值进行关联，合并成一张表。
因为须要应用 file_name 作为 Hive 表二级动静分区，也防止大量数据 join 导致 OOM，这里按 dir_name 别离进行 join 操作，不同的分区批改对应的 dir_name 即可。

join_sql = """insert overwrite table laion5b.dmd_image_path_suffix_aesthetics_join_view PARTITION (dir_name ='laion2B-en', file_name)
select A.sample_id,
    A.`hash`,
    B.image_path,
    B.image_suffix,
    A.license,
    A.nsfw,
    C.aesthetic_score,
    A.file_name
from laion5b.ods_parquet_short_view A
    left join laion5b.ods_image_path_view B on A.`hash` = B.`hash`
    and B.dir_name = 'laion2B-en'
    left join laion5b.ods_improved_aesthetics_parquet_view C on A.`hash` = C.`hash`
where A.dir_name = 'laion2B-en'
"""

join_df = sqlContext.sql(join_sql)

表 5 与表 6 通过 hash 字段进行 join，失去 result 后果表 7。

sc = spark.sparkContext
sqlContext = SQLContext(sc)

join_sql = """insert overwrite table laion5b.dms_parquet_result_view PARTITION (dir_name ='laion2B-en', file_name)
select B.id, B.sample_id, A.url, A.text, A.width, A.height, 
    A.similarity, A.`hash`, A.punsafe, A.pwatermark,
    if(B.image_suffix is null or B.image_suffix = 'jpg', B.image_path, 
    regexp_replace(B.image_path, right(B.image_path, 3), B.image_suffix)) as image_path,
    B.image_suffix, B.LICENSE, B.NSFW, B.aesthetic_score, A.language, B.file_name
from laion5b.ods_joined_parquet_view A
    left join (
        select id, sample_id, image_path, image_suffix, LICENSE, 
            NSFW, aesthetic_score, `hash`, dir_name, file_name, 
            ROW_NUMBER() OVER (PARTITION BY `hash` order by `hash`) as rn
        from laion5b.dmd_image_path_suffix_aesthetics_join_view
        where dir_name = 'laion2B-en'
    ) B on A.dir_name = B.dir_name
    and A.`hash` = B.`hash`
    and B.rn = 1
where A.dir_name = 'laion2B-en'
"""

join_df = sqlContext.sql(join_sql)

4-3.Data write

最初将 Hive 后果表导出为 snappy 压缩格局的 parquet 文件，再上传到对象存储就能够应用了。

sc = spark.sparkContext
sql_context = SQLContext(sc)


write_df = sql_context.sql("""
select sample_id, url, text, width, height, similarity, hash, 
    punsafe, pwatermark, image_suffix, license, nsfw, aesthetic_score, language, file_name 
    from laion5b.parquet_improve_result_view 
    where dir_name='laion2B-multi'
""")

write_df.repartition(1).write.parquet("/datasets/result/multi/", mode="overwrite", partitionBy='file_name', compression="snappy")

五、LAION-5B 媒体图片下载

在 OpenDataLab 网站下载到 parquet 文件后，能够依据这份元数据下载对应的图片文件。
当初，咱们也开源了 LAION-5B 图片下载代码，github 开源地址如下：
https://github.com/opendatalab/laion5b-downloader

耗时 25 天，目前下载的图片总量为 5065377962 张（因 url 链接和网站起因，局部图片无奈下载），总存储量为 300+TB。

六、parquet 不可不知（附下载链接）

在对官网 parquet 进行解决时，发现了数据中的几个小问题，也同步你知。

1、初始 LAION-5B parquet 文件中 sample id 字段反复度较高，有同学应用 sample_id 字段作为惟一索引，为了防止数据问题，倡议应用 hash 替换。

咱们来看看 sample_id 的反复状况：

当然，hash 在按 laion2b-en、laion2b-multi、独立分类的子集中共有 6 条反复值，laion1b-nolang 无数据反复。
在应用时留神下即可，hash 值反复状况如下表：

2、3 个分类子集的图片准确数量如下，能够用来比照图片是否有缺失问题。

最初，有须要的同学能够在 OpenDataLab 下载解决好的 parquet 文件。
OpenDataLab parquet 文件下载链接：
https://opendatalab.com/LAION-5B

-END-
作者 | 喻佳、张文坚