时间戳

首页 > 解决方案 > 我可以仅使用 .pth 文件复制和运行其他人的模型吗？

python - 我可以仅使用 .pth 文件复制和运行其他人的模型吗？

问题描述

我试图准确了解 .pth 文件是什么以及如何使用它们来测试和运行其他人的模型。

例如，查看这样的存储库：https ://github.com/TencentYoutuResearch/PedestrianDetection-NohNMS

我是否可以下载 .pth 文件并从中创建一个执行/计算相同的 Keras 模型？

我在想而不是下载并尝试编译上面的整个存储库，如果我只想测试他们的模型，也许我可以下载他们的 .pth 文件并执行类似“my_model = keras.model(pth_file)”之类的操作，你会得到这个想法。

标签： pythonkeras

解决方案

似乎https://github.com/gmalivenko/pytorch2keras库可以满足您的需求。

在自述文件中，他们也有一个例子来做到这一点。

它是 PyTorch 图到 Keras（Tensorflow 后端）模型的转换器。首先，我们需要加载（或创建）一个有效的 PyTorch 模型：

class TestConv2d(nn.Module):
    """
    Module for Conv2d testing
    """

    def __init__(self, inp=10, out=16, kernel_size=3):
        super(TestConv2d, self).__init__()
        self.conv2d = nn.Conv2d(inp, out, stride=1, kernel_size=kernel_size, bias=True)

    def forward(self, x):
        x = self.conv2d(x)
        return x

model = TestConv2d()

# load weights here
# model.load_state_dict(torch.load(path_to_weights.pth))

下一步 - 创建一个具有正确形状的虚拟变量：

input_np = np.random.uniform(0, 1, (1, 10, 32, 32))
input_var = Variable(torch.FloatTensor(input_np))

我们使用虚拟变量来跟踪模型（使用 jit.trace）：

from pytorch2keras import pytorch_to_keras
# we should specify shape of the input tensor
k_model = pytorch_to_keras(model, input_var, [(10, 32, 32,)], verbose=True)

您还可以将 H 和 W 尺寸设置为 None 以使您的模型与形状无关（例如，完全卷积网络）：

from pytorch2keras.converter import pytorch_to_keras
# we should specify shape of the input tensor
k_model = pytorch_to_keras(model, input_var, [(10, None, None,)], verbose=True)

就这样！如果所有模块都正确转换，Keras 模型将存储在 k_model 变量中。

但这只有在所有层都受支持的情况下才有效，您可以在文档中找到列表。

推荐阅读