Default optimize_for_mobile() on TorchScript models (#2908)
Per https://pytorch.org/tutorials/recipes/script_optimized.html this should improve performance on torchscript models (and maybe coreml models also since coremltools operates on a torchscript model input, though this still requires testing).
This commit is contained in:
1 parent
f662c18507
commit
28db237639
1 file changed
+2
@@ -12,6 +12,7 @@ sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
|
from torch.utils.mobile_optimizer import optimize_for_mobile
|
||||||
|
|
||||||
import models
|
import models
|
||||||
from models.experimental import attempt_load
|
from models.experimental import attempt_load
|
||||||
@@ -65,6 +66,7 @@ if __name__ == '__main__':
|
|||||||
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
||||||
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
||||||
ts = torch.jit.trace(model, img, strict=False)
|
ts = torch.jit.trace(model, img, strict=False)
|
||||||
|
ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html
|
||||||
ts.save(f)
|
ts.save(f)
|
||||||
print(f'{prefix} export success, saved as {f}')
|
print(f'{prefix} export success, saved as {f}')
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
Reference in new issue
Block a user