将 pytorch_model.bin 拆分为块后无法从 pytorch 检查点加载权重答案

【问题标题】：Unable to load weights from pytorch checkpoint after splitting pytorch_model.bin into chunks将 pytorch_model.bin 拆分为块后无法从 pytorch 检查点加载权重
【发布时间】：2021-02-02 07:05:27
【问题描述】：

我需要传输预训练的 deeppavlov ruBERT 模型的 pytorch_model.bin，但我有文件大小限制。所以我用python把它分成块，以正确的顺序传输和重新组装。但是，文件的大小增加了，当我尝试使用 BertModel.from_pretrained(pytorch_model.bin) 加载生成的文件时，我收到了一个错误：

During handling of the above exception, another exception occurred:
OSError: Unable to load weights from pytorch checkpoint <...>

所以我的问题是：真的可以这样拆分文件吗？我在拆分和重组文件的方式上可能有错误。但是，这也可能是某些版本不匹配。

我获取块的python代码：

chunk_size = 40000000
file_num = 1
with open("pytorch_model.bin", "rb") as f:
    chunk = f.read(chunk_size)
    while chunk:
        with open("chunk_" + str(file_num), "wb") as chunk_file:
            chunk_file.write(chunk)
        file_num += 1
        chunk = f.read(chunk_size)

重组一个文件的代码：

chunks = !ls | grep chunk_
chunks = sorted(chunks, key=lambda x: int(x.split("_")[-1]))

for chunk in chunks:
    with open(chunk, "rb") as f:
        contents = f.read()
    if chunk == chunks[0]:
        write_mode = "wb"
    else:
        write_mode = "ab"
    with open("pytorch_model.bin", write_mode) as f:
        f.write(contents)

python 3.7.0，火炬 1.5.1，变形金刚 4.2.2。我无法移动大于 40 MB 的文件。

TIA 为您提供帮助！

【问题讨论】：

标签： python serialization pytorch huggingface-transformers bin

【解决方案1】：

我建议不要使用 python。

使用命令行split 和cat 拆分大文件，并将拆分后的文件合并到另一侧的单个文件中（this thread 显示如何）。

我建议你使用md5sum（或其他校验和功能）来验证你在接收端组装的pytorch_model.bin文件确实与原始文件相同。

【讨论】：

【解决方案2】：

我与我的团队核实了保存模型时使用的转换器和 pytorch 的版本。它与我用来加载模型的版本不同。所以我安装了保存模型时使用的版本，然后重新尝试加载。它奏效了。

【讨论】：