From 7393b26e633ef568ebe134cbd0e2a04be7c374ee Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 24 Apr 2021 01:31:11 +0200
Subject: [PATCH] Add file_size() function (#2911)

* Add file_size() function

* Update export.py
---
 models/export.py | 9 +++++----
 utils/general.py | 5 +++++
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/models/export.py b/models/export.py
index a820ff94ac17..38fefa7e891c 100644
--- a/models/export.py
+++ b/models/export.py
@@ -17,7 +17,7 @@
 import models
 from models.experimental import attempt_load
 from utils.activations import Hardswish, SiLU
-from utils.general import colorstr, check_img_size, check_requirements, set_logging
+from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
 from utils.torch_utils import select_device
 
 if __name__ == '__main__':
@@ -60,6 +60,7 @@
     model.model[-1].export = not opt.grid  # set Detect() layer grid export
     for _ in range(2):
         y = model(img)  # dry runs
+    print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
 
     # TorchScript export -----------------------------------------------------------------------------------------------
     prefix = colorstr('TorchScript:')
@@ -69,7 +70,7 @@
         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}')
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
         print(f'{prefix} export failure: {e}')
 
@@ -103,7 +104,7 @@
                 onnx.save(model_onnx, f)
             except Exception as e:
                 print(f'{prefix} simplifier failure: {e}')
-        print(f'{prefix} export success, saved as {f}')
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
         print(f'{prefix} export failure: {e}')
 
@@ -117,7 +118,7 @@
         model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
         f = opt.weights.replace('.pt', '.mlmodel')  # filename
         model.save(f)
-        print(f'{prefix} export success, saved as {f}')
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
         print(f'{prefix} export failure: {e}')
 
diff --git a/utils/general.py b/utils/general.py
index 92c6ca5df208..ba88759c2983 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -61,6 +61,11 @@ def emojis(str=''):
     return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
 
 
+def file_size(file):
+    # Return file size in MB
+    return Path(file).stat().st_size / 1e6
+
+
 def check_online():
     # Check internet connectivity
     import socket