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).pull/2909/head
parent
f662c18507
commit
28db237639
|
@ -12,6 +12,7 @@ sys.path.append('./') # to run '$ python *.py' files in subdirectories
|
|||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.mobile_optimizer import optimize_for_mobile
|
||||
|
||||
import models
|
||||
from models.experimental import attempt_load
|
||||
|
@ -65,6 +66,7 @@ if __name__ == '__main__':
|
|||
print(f'\n{prefix} starting export with torch {torch.__version__}...')
|
||||
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
||||
ts = torch.jit.trace(model, img, strict=False)
|
||||
ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html
|
||||
ts.save(f)
|
||||
print(f'{prefix} export success, saved as {f}')
|
||||
except Exception as e:
|
||||
|
|
Loading…
Reference in New Issue