Question

我尝试过的一个选项是腌制vocab，并使用Extrafiles arg保存

import torch
import pickle

class Vocab(object):
    pass

vocab = Vocab()
pickle.dump(open('path/to/vocab.pkl','w'))

m = torch.jit.ScriptModule()

## I am not sure about the usage of this arg, the docs didn't help me
extra_files = torch._C.ExtraFilesMap()
extra_files['vocab.pkl'] = 'path/to/vocab.pkl'
# I also tried  pickle.dumps(vocab), and directly vocab

torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)

## Load with extra files.
files = {'vocab.pkl': ''}
torch.jit.load('scriptmodule.pt', _extra_files = files)

这给

TypeError: import_ir_module(): incompatible function arguments. The following argument types are supported:
    1. (arg0: Callable[[List[str]], torch._C.ScriptModule], arg1: str, arg2: object, arg3: torch._C.ExtraFilesMap) -> None

显然，另一个选择是单独加载泡菜，但是我正在寻找单个文件的选择。

如果可以将vocab添加到火炬手中，那将是一件很不错的事情……也很高兴知道是否有某些我不知道的不这样做的原因。

Answer 1

我认为torch.jit.load的文档不正确。您需要创建一个ExtraFilesmap（）对象来加载保存的文件。

以下是我如何工作的示例：第1步：保存模型

extra_files = torch._C.ExtraFilesMap()
extra_files['foo.txt'] = 'bar'
traced_script_module.save(serialized_model_path, _extra_files=extra_files)

第2步：加载模型

files = torch._C.ExtraFilesMap()
files['foo.txt'] = ''
loaded_model = torch.jit.load(serialized_model_path, _extra_files=files)
print(files)

Answer 2

问题位于torch.jit.load中。尝试检查您的map_location

Answer 3

假设vocab是受支持的类型，则可以将其作为TorchScript attribute添加到模型中，以将其与模型一起存储在1个文件中（这样就不必处理{{1 }}。

然后您的加载代码变为

_extra_files

在torch.jit.save中_extra_files arg的正确用法是什么

3 个答案: