|
|
@@ -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: |