如何将Kubric实现的MOVi数据集从GCS下载至本地磁盘？

针对你遇到的tfds.load无法触发下载、ds.save内存溢出问题，以下是三个可行的解决办法：

方法一：用gsutil直接同步GCS数据集文件（推荐） #

MOVi数据集已经在GCS上以TFDS格式存储，直接将对应文件夹同步到本地是最高效的方式，完全避免内存问题：

1. 确保已安装gsutil（属于Google Cloud SDK组件）
2. 执行同步命令，替换/path/to/local/tfds_dir为你要存储数据集的本地目录：

   gsutil -m rsync -r gs://kubric-public/tfds/movi_b /path/to/local/tfds_dir/movi_b
   

   • -m：启用多线程加速传输
   • rsync：增量同步，仅下载缺失或更新的文件，支持断点续传
3. 后续加载本地数据时，指定本地data_dir即可：

   ds, info = tfds.load("movi_b", data_dir="/path/to/local/tfds_dir", with_info=True)
   

方法二：分批处理数据集，避免内存溢出 #

如果不想使用gsutil，可以修改ds.save的逻辑，逐个样本分块保存，彻底解决内存占用问题：

import tensorflow as tf
import tensorflow_datasets as tfds

# 流式加载数据集
ds = tfds.load("movi_b", data_dir="gs://kubric-public/tfds", split="train")

# 分块保存每个样本
save_root = "./local_movi_b/train"
for idx, sample in enumerate(ds):
    # 每个分片仅保存单个样本
    tf.data.Dataset.from_tensors(sample).save(f"{save_root}/shard_{idx:06d}")
    # 打印进度（可选）
    if idx % 100 == 0:
        print(f"Completed {idx} samples")

# 后续加载本地数据的方式
loaded_ds = tf.data.Dataset.list_files(f"{save_root}/shard_*")\
                           .interleave(tf.data.Dataset.load, num_parallel_calls=tf.data.AUTOTUNE)

方法三：强制TFDS重新下载并准备数据集 #

由于TFDS默认会优先使用GCS上的已预处理数据集，导致download=True无效，可强制触发本地预处理流程：

import tensorflow_datasets as tfds

# 指定本地存储目录
local_data_dir = "./local_tfds_dir"
builder = tfds.builder("movi_b", data_dir=local_data_dir)

# 强制跳过GCS缓存，重新下载预处理
builder.download_and_prepare(
    download_dir="./temp_download",  # 临时存储原始下载文件
    download_config=tfds.download.DownloadConfig(
        try_download_gcs=False,  # 禁用GCS已准备数据的自动使用
    )
)

# 加载本地数据集
ds = builder.as_dataset(split="train")

注意事项 #

• MOVi系列数据集体积较大（如movi_b约几十GB），确保本地磁盘有足够剩余空间
• gsutil同步时，网络中断后重新执行命令即可续传
• 分块保存的数据集可以通过tf.data.Dataset.interleave高效合并加载

内容的提问来源于stack exchange，提问作者FooBar