如何将PyTorch中的torchscript模型转换为普通的nn.Module？

Question

我正在以以下方式加载torchscript模型：

model = torch.jit.load("model.pt").to(device)

该模型的子模块被识别为RecursiveScriptModule。我想整理上传的重量，为了使它更简单，并将它们投到torch.float32，最好将所有这些东西转换成普通的PyTorch nn.Module。

在正式的docs https://pytorch.org/docs/stable/jit.html中，我们被告知如何将nn.Module转换为torchscript，但是我还没有发现在相反方向上这样做的例子。有办法这样做吗？

本文给出了加载模型预训练模型的实例：https://github.com/openai/CLIP/blob/main/notebooks/Interacting_with_CLIP.ipynb

Answer 1

您可以尝试加载它，例如state_dict = torch.load(src).state_dict()。然后手动转换每个键和值new_v = state_dict[k].cpu().float()。