是否可以在不首先在本地持久化的情况下从 GCS 存储桶 URL 加载预训练的 Pytorch 模型？

Question

我是在 Google Dataflow 的背景下问这个问题的，但也是一般的。

使用 PyTorch，我可以引用包含多个文件的本地目录，这些文件构成一个预训练模型。我碰巧使用的是 Roberta 模型，但其他人的界面是一样的。

ls some-directory/
      added_tokens.json
      config.json             
      merges.txt              
      pytorch_model.bin       
      special_tokens_map.json vocab.json

ls some-directory/
      added_tokens.json
      config.json             
      merges.txt              
      pytorch_model.bin       
      special_tokens_map.json vocab.json

但是，我的预训练模型存储在 GCS 存储桶中。让我们称之为gs://my-bucket/roberta/。

在 Google Dataflow 中加载这个模型的上下文中，我试图保持无状态并避免持久化到磁盘，所以我更喜欢直接从 GCS 获取这个模型。据我了解，PyTorch 通用接口方法from_pretrained()可以采用本地目录或 URL 的字符串表示形式。但是，我似乎无法从 GCS URL 加载模型。

from pytorch_transformers import RobertaModel

# this works
model = RobertaModel.from_pretrained('/path/to/some-directory/')

如果我尝试使用目录 blob 的公共 https URL，它也会失败，尽管这可能是由于缺乏身份验证，因为在可以创建客户端的 python 环境中引用的凭据不会转换为公共请求https://storage.googleapis

# this fails, probably due to auth
bucket = gcs_client.get_bucket('my-bucket')
directory_blob = bucket.blob(prefix='roberta')
model = RobertaModel.from_pretrained(directory_blob.public_url)
# ValueError: No JSON object could be decoded

# and for good measure, it also fails if I append a trailing /
model = RobertaModel.from_pretrained(directory_blob.public_url + '/')
# ValueError: No JSON object could be decoded

我知道GCS 实际上没有子目录，它实际上只是存储桶名称下的一个平面命名空间。但是，似乎我被身份验证的必要性和 PyTorch 阻止了不说话gs://。

我可以通过首先在本地保存文件来解决这个问题。

# this fails
model = RobertaModel.from_pretrained('gs://my-bucket/roberta/')
# ValueError: unable to parse gs://mahmed_bucket/roberta-base as a URL or as a local path

但这似乎是一个黑客，我一直在想我一定错过了一些东西。当然有一种方法可以保持无状态，而不必依赖磁盘持久性！

• 那么有没有办法加载存储在 GCS 中的预训练模型？
• 在此上下文中执行公共 URL 请求时，是否可以进行身份​​验证？
• 即使有一种方法可以进行身份​​验证，子目录不存在仍然是一个问题吗？

谢谢您的帮助！我也很高兴被指出任何重复的问题，因为我肯定找不到任何问题。

---

编辑和澄清

• 我的 Python 会话已经通过了 GCS 身份验证，这就是为什么我能够在本地下载 blob 文件，然后使用load_frompretrained()

• load_frompretrained() 需要目录引用，因为它需要问题顶部列出的所有文件，而不仅仅是 pytorch-model.bin

• 为了澄清问题 #2，我想知道是否有某种方法可以为 PyTorch 方法提供一个嵌入了加密凭据或类似内容的请求 URL。有点远，但我想确保我没有错过任何东西。

• 为了澄清问题 #3（除了对下面一个答案的评论），即使有一种方法可以在我不知道的 URL 中嵌入凭据，我仍然需要引用一个目录而不是单个 blob，并且我不知道 GCS 子目录是否会被识别为这样，因为（正如 Google 文档所述）GCS 中的子目录是一种错觉，它们并不代表真正的目录结构。所以我认为这个问题无关紧要，或者至少被问题 #2 阻止，但这是我追寻的一个话题，所以我仍然很好奇。

Answer 1

主要编辑：

您可以在 Dataflow Worker 上安装 Wheel 文件，还可以使用 Worker 临时存储在本地保存二进制文件！

确实（目前截至 2019 年 11 月）您无法通过提供--requirements参数来做到这一点。相反，你必须setup.py像这样使用。假设大写字母中的任何常量在其他地方定义。

REQUIRED_PACKAGES = [
    'torch==1.3.0',
    'pytorch-transformers==1.2.0',
]

setup(
    name='project_dir',
    version=VERSION,
    packages=find_packages(),
    install_requires=REQUIRED_PACKAGES)

运行脚本

python setup.py sdist

python project_dir/my_dataflow_job.py \
--runner DataflowRunner \
--project ${GCP_PROJECT} \
--extra_package dist/project_dir-0.1.0.tar.gz \
# SNIP custom args for your job and required Dataflow Temp and Staging buckets #

在作业中，这里是在自定义数据流运算符的上下文中下载并使用 GCS 的模型。为了方便起见，我们将一些实用方法包装在一个单独的模块中（对于解决数据流依赖项上传很重要），并将它们导入到自定义运算符的本地范围，而不是全局范围。

class AddColumn(beam.DoFn):
    PRETRAINED_MODEL = 'gs://my-bucket/blah/roberta-model-files'

    def get_model_tokenizer_wrapper(self):
        import shutil
        import tempfile
        import dataflow_util as util
        try:
            return self.model_tokenizer_wrapper
        except AttributeError:
            tmp_dir = tempfile.mkdtemp() + '/'
            util.download_tree(self.PRETRAINED_MODEL, tmp_dir)
            model, tokenizer = util.create_model_and_tokenizer(tmp_dir)
            model_tokenizer_wrapper = util.PretrainedPyTorchModelWrapper(
                model, tokenizer)
            shutil.rmtree(tmp_dir)
            self.model_tokenizer_wrapper = model_tokenizer_wrapper
            logging.info(
                'Successfully created PretrainedPyTorchModelWrapper')
            return self.model_tokenizer_wrapper

    def process(self, elem):
        model_tokenizer_wrapper = self.get_model_tokenizer_wrapper()

        # And now use that wrapper to process your elem however you need.
        # Note that when you read from BQ your elements are dictionaries
        # of the column names and values for each BQ row.

代码库中单独模块中的实用程序函数。在我们的项目根目录中，它位于 dataflow_util/init.py 中，但您不必这样做。

from contextlib import closing
import logging

import apache_beam as beam
import numpy as np
from pytorch_transformers import RobertaModel, RobertaTokenizer
import torch

class PretrainedPyTorchModelWrapper():
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

def download_tree(gcs_dir, local_dir):
    gcs = beam.io.gcp.gcsio.GcsIO()
    assert gcs_dir.endswith('/')
    assert local_dir.endswith('/')
    for entry in gcs.list_prefix(gcs_dir):
        download_file(gcs, gcs_dir, local_dir, entry)


def download_file(gcs, gcs_dir, local_dir, entry):
    rel_path = entry[len(gcs_dir):]
    dest_path = local_dir + rel_path
    logging.info('Downloading %s', dest_path)
    with closing(gcs.open(entry)) as f_read:
        with open(dest_path, 'wb') as f_write:
            # Download the file in chunks to avoid requiring large amounts of
            # RAM when downloading large files.
            while True:
                file_data_chunk = f_read.read(
                    beam.io.gcp.gcsio.DEFAULT_READ_BUFFER_SIZE)
                if len(file_data_chunk):
                    f_write.write(file_data_chunk)
                else:
                    break


def create_model_and_tokenizer(local_model_path_str):
    """
    Instantiate transformer model and tokenizer

      :param local_model_path_str: string representation of the local path 
             to the directory containing the pretrained model
      :return: model, tokenizer
    """
    model_class, tokenizer_class = (RobertaModel, RobertaTokenizer)

    # Load the pretrained tokenizer and model
    tokenizer = tokenizer_class.from_pretrained(local_model_path_str)
    model = model_class.from_pretrained(local_model_path_str)

    return model, tokenizer

伙计们，你已经得到了它！更多详细信息可以在这里找到：https ://beam.apache.org/documentation/sdks/python-pipeline-dependency/

我发现这整个问题链是无关紧要的，因为 Dataflow 只允许您在工作人员上安装源分发包，这意味着您实际上无法安装 PyTorch。

当您提供requirements.txt文件时，Dataflow 将使用--no-binary标记进行安装，该标记阻止安装 Wheel (.whl) 软件包并仅允许源分发 (.tar.gz)。我决定尝试在 Google Dataflow 上推出我自己的 PyTorch 源代码发行版，其中一半是 C++，一半是 Cuda，另一部分谁知道这是一个傻瓜的差事。

感谢大家一路以来的投入。