_extra_files arg 在 torch.jit.save 中的正确用法是什么
what is the right usage of _extra_files arg in torch.jit.save
我尝试过的一个选项是 pickling 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
其他选项显然是单独加载 pickle,但我正在寻找单个文件选项。
如果可以将 vocab 添加到 torchscript 中就好了……如果知道是否有某些我显然不知道的不这样做的原因也很好。
问题位于torch.jit.load。
尝试检查您的 map_location
假设 vocab
是受支持的类型,您可以将其作为 TorchScript attribute 添加到模型中,以便将其与模型一起存储在 1 个文件中(这样您就不必处理 _extra_files
).
那么你的加载代码就变成了
torch.jit.load('scriptmodule.pt')
我认为 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)
我尝试过的一个选项是 pickling 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
其他选项显然是单独加载 pickle,但我正在寻找单个文件选项。
如果可以将 vocab 添加到 torchscript 中就好了……如果知道是否有某些我显然不知道的不这样做的原因也很好。
问题位于torch.jit.load。 尝试检查您的 map_location
假设 vocab
是受支持的类型,您可以将其作为 TorchScript attribute 添加到模型中,以便将其与模型一起存储在 1 个文件中(这样您就不必处理 _extra_files
).
那么你的加载代码就变成了
torch.jit.load('scriptmodule.pt')
我认为 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)