pytorch torch.jit.trace()使用模型转换要注意的问题

https://www.jianshu.com/p/8206b198533d如果代码中有if条件控制，尽量避免使用torch.jit.trace来转换代码，因为它不能处理变化条件，如果非要用trace的话，可以把if条件控制改成别的形式，比如：def f(x):if x > 0:return Falseelse:return True可以改成:def f(x):return x <=

苏三福

18325人浏览 · 2020-09-28 14:15:26

苏三福 · 2020-09-28 14:15:26 发布

https://www.jianshu.com/p/8206b198533d

如果代码中有if条件控制，尽量避免使用torch.jit.trace来转换代码，因为它不能处理变化条件，如果非要用trace的话，可以把if条件控制改成别的形式，比如：

def f(x):
  if x > 0:
    return False
  else:
    return True

可以改成:

def f(x):
  return x <= 0

jit不能转换第三方Python库中的函数，尽量所有代码都使用pytorch实现，如果速度不理想的话，可以参考github上的pytorch/extension-script项目，用C++实现需要的功能，然后注册成jit操作，最后转成torchscript。
如果要转Mobilenet，最好使用pytorch1.3以上，否则识别不出来其中的depth wise conv，转换出来的torchscript模型会比原模型大很多。
模型的forward函数中尽量不要包含中文注释。

作者：dalalaa
链接：https://www.jianshu.com/p/8206b198533d
来源：简书
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。