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

问题描述

我需要传输预训练的 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。

使用命令行 split 和 cat 拆分大文件并将拆分的文件在另一侧合并为一个文件（this thread 展示了如何操作）。

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

,

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