- 如果代码中有
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
函数中尽量不要包含中文注释。