2

只是想澄清我对 JIT 和 TorchScripts 工作方式的理解并澄清一个特定的例子。

因此，如果我没有错torch.jit.script，将我的方法或模块转换为 TorchScript。我可以在 python 之外的环境中使用我的 TorchScript 编译模块，但也可以在 python 中使用它并进行预期的改进和优化。一个类似的情况torch.jit.trace是跟踪权重和操作，但大致遵循类似的想法。

如果是这种情况，TorchScripted 模块通常应该至少与 python 解释器的典型推理时间一样快。在进行了一些实验后，我观察到它通常比典型的解释器推理时间要慢，并且在阅读后发现显然 TorchScripted 模块需要“预热”一下，以实现其最佳性能。这样做时，我看到推理时间没有任何变化，它变得更好但不足以称为对典型的做事方式（python解释器）的改进。此外，我使用了一个名为 的第三方库torch_tvm，如果启用该库，据说可以将任何方式的 jit-ing 模块的推理时间减半。

直到现在这一切都没有发生，我真的无法说出原因。

以下是我的示例代码，以防我做错了什么 -

class TrialC(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(1024, 2048)
        self.l2 = nn.Linear(2048, 4096)
        self.l3 = nn.Linear(4096, 4096)
        self.l4 = nn.Linear(4096, 2048)
        self.l5 = nn.Linear(2048, 1024)

    def forward(self, input):
        out = self.l1(input)
        out = self.l2(out)
        out = self.l3(out)
        out = self.l4(out)
        out = self.l5(out)
        return out 

if __name__ == '__main__':
    # Trial inference input 
    TrialC_input = torch.randn(1, 1024)
    warmup = 10

    # Record time for typical inference 
    model = TrialC()
    start = time.time()
    model_out = model(TrialC_input)
    elapsed = time.time() - start 

    # Record the 10th inference time (10 warmup) for the optimized model in TorchScript 
    script_model = torch.jit.script(TrialC())
    for i in range(warmup):
        start_2 = time.time()
        model_out_check_2 = script_model(TrialC_input)
        elapsed_2 = time.time() - start_2

    # Record the 10th inference time (10 warmup) for the optimized model in TorchScript + TVM optimization
    torch_tvm.enable()
    script_model_2 = torch.jit.trace(TrialC(), torch.randn(1, 1024))
    for i in range(warmup):
        start_3 = time.time()
        model_out_check_3 = script_model_2(TrialC_input)
        elapsed_3 = time.time() - start_3 
    
    print("Regular model inference time: {}s\nJIT compiler inference time: {}s\nJIT Compiler with TVM: {}s".format(elapsed, elapsed_2, elapsed_3))

以下是上述代码在我的 CPU 上的结果 -

Regular model inference time: 0.10335588455200195s
JIT compiler inference time: 0.11449170112609863s
JIT Compiler with TVM: 0.10834860801696777s

对此的任何帮助或澄清将不胜感激！

4

0 回答 0