Commit
•
1b1ab4c
1
Parent(s):
646386f
Add file_size() function (#2911)
Browse files* Add file_size() function
* Update export.py
- models/export.py +5 -4
- utils/general.py +5 -0
models/export.py
CHANGED
@@ -17,7 +17,7 @@ from torch.utils.mobile_optimizer import optimize_for_mobile
|
|
17 |
import models
|
18 |
from models.experimental import attempt_load
|
19 |
from utils.activations import Hardswish, SiLU
|
20 |
-
from utils.general import colorstr, check_img_size, check_requirements, set_logging
|
21 |
from utils.torch_utils import select_device
|
22 |
|
23 |
if __name__ == '__main__':
|
@@ -60,6 +60,7 @@ if __name__ == '__main__':
|
|
60 |
model.model[-1].export = not opt.grid # set Detect() layer grid export
|
61 |
for _ in range(2):
|
62 |
y = model(img) # dry runs
|
|
|
63 |
|
64 |
# TorchScript export -----------------------------------------------------------------------------------------------
|
65 |
prefix = colorstr('TorchScript:')
|
@@ -69,7 +70,7 @@ if __name__ == '__main__':
|
|
69 |
ts = torch.jit.trace(model, img, strict=False)
|
70 |
ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html
|
71 |
ts.save(f)
|
72 |
-
print(f'{prefix} export success, saved as {f}')
|
73 |
except Exception as e:
|
74 |
print(f'{prefix} export failure: {e}')
|
75 |
|
@@ -103,7 +104,7 @@ if __name__ == '__main__':
|
|
103 |
onnx.save(model_onnx, f)
|
104 |
except Exception as e:
|
105 |
print(f'{prefix} simplifier failure: {e}')
|
106 |
-
print(f'{prefix} export success, saved as {f}')
|
107 |
except Exception as e:
|
108 |
print(f'{prefix} export failure: {e}')
|
109 |
|
@@ -117,7 +118,7 @@ if __name__ == '__main__':
|
|
117 |
model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
118 |
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
119 |
model.save(f)
|
120 |
-
print(f'{prefix} export success, saved as {f}')
|
121 |
except Exception as e:
|
122 |
print(f'{prefix} export failure: {e}')
|
123 |
|
|
|
17 |
import models
|
18 |
from models.experimental import attempt_load
|
19 |
from utils.activations import Hardswish, SiLU
|
20 |
+
from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
|
21 |
from utils.torch_utils import select_device
|
22 |
|
23 |
if __name__ == '__main__':
|
|
|
60 |
model.model[-1].export = not opt.grid # set Detect() layer grid export
|
61 |
for _ in range(2):
|
62 |
y = model(img) # dry runs
|
63 |
+
print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
|
64 |
|
65 |
# TorchScript export -----------------------------------------------------------------------------------------------
|
66 |
prefix = colorstr('TorchScript:')
|
|
|
70 |
ts = torch.jit.trace(model, img, strict=False)
|
71 |
ts = optimize_for_mobile(ts) # https://pytorch.org/tutorials/recipes/script_optimized.html
|
72 |
ts.save(f)
|
73 |
+
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
74 |
except Exception as e:
|
75 |
print(f'{prefix} export failure: {e}')
|
76 |
|
|
|
104 |
onnx.save(model_onnx, f)
|
105 |
except Exception as e:
|
106 |
print(f'{prefix} simplifier failure: {e}')
|
107 |
+
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
108 |
except Exception as e:
|
109 |
print(f'{prefix} export failure: {e}')
|
110 |
|
|
|
118 |
model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
|
119 |
f = opt.weights.replace('.pt', '.mlmodel') # filename
|
120 |
model.save(f)
|
121 |
+
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
122 |
except Exception as e:
|
123 |
print(f'{prefix} export failure: {e}')
|
124 |
|
utils/general.py
CHANGED
@@ -61,6 +61,11 @@ def emojis(str=''):
|
|
61 |
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
62 |
|
63 |
|
|
|
|
|
|
|
|
|
|
|
64 |
def check_online():
|
65 |
# Check internet connectivity
|
66 |
import socket
|
|
|
61 |
return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
|
62 |
|
63 |
|
64 |
+
def file_size(file):
|
65 |
+
# Return file size in MB
|
66 |
+
return Path(file).stat().st_size / 1e6
|
67 |
+
|
68 |
+
|
69 |
def check_online():
|
70 |
# Check internet connectivity
|
71 |
import socket
|