Default optimize_for_mobile() on TorchScript models (#2908)
Browse filesPer 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).
- models/export.py +2 -0
models/export.py
CHANGED
@@ -12,6 +12,7 @@ sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
|
12 |
|
13 |
import torch
|
14 |
import torch.nn as nn
|
|
|
15 |
|
16 |
import models
|
17 |
from models.experimental import attempt_load
|
@@ -65,6 +66,7 @@ if __name__ == '__main__':
|
|
65 |
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
66 |
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
67 |
ts = torch.jit.trace(model, img, strict=False)
|
|
|
68 |
ts.save(f)
|
69 |
print(f'{prefix} export success, saved as {f}')
|
70 |
except Exception as e:
|
|
|
12 |
|
13 |
import torch
|
14 |
import torch.nn as nn
|
15 |
+
from torch.utils.mobile_optimizer import optimize_for_mobile
|
16 |
|
17 |
import models
|
18 |
from models.experimental import attempt_load
|
|
|
66 |
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
67 |
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
68 |
ts = torch.jit.trace(model, img, strict=False)
|
69 |
+
ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html
|
70 |
ts.save(f)
|
71 |
print(f'{prefix} export success, saved as {f}')
|
72 |
except Exception as e:
|