Add git info to training checkpoints (#9655)

* Add git status on train checkpoints * Update * Update * Update * Update * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update general.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> * Update general.py Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
ultralytics · Nov 19, 2022 · 9286336 · 9286336
1 parent 9bc6034
commit 9286336
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 2 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -2,6 +2,7 @@
 # Usage: pip install -r requirements.txt
 
 # Base ------------------------------------------------------------------------
+gitpython
 ipython  # interactive notebook
 matplotlib>=3.2.2
 numpy>=1.18.5

diff --git a/train.py b/train.py
@@ -47,7 +47,7 @@
 from utils.callbacks import Callbacks
 from utils.dataloaders import create_dataloader
 from utils.downloads import attempt_download, is_url
-from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_status,
+from utils.general import (GIT, LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_status,
                            check_img_size, check_requirements, check_suffix, check_yaml, colorstr, get_latest_run,
                            increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
                            labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer,
@@ -376,6 +376,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                     'updates': ema.updates,
                     'optimizer': optimizer.state_dict(),
                     'opt': vars(opt),
+                    'git': GIT,
                     'date': datetime.now().isoformat()}
 
                 # Save last, best and delete

diff --git a/utils/general.py b/utils/general.py
@@ -29,6 +29,7 @@
 from zipfile import ZipFile, is_zipfile
 
 import cv2
+import git
 import IPython
 import numpy as np
 import pandas as pd
@@ -344,6 +345,22 @@ def check_git_status(repo='ultralytics/yolov5', branch='master'):
     LOGGER.info(s)
 
 
+@WorkingDirectory(ROOT)
+def check_git(path='.'):
+    # YOLOv5 git check, return git {remote, branch, commit}
+    try:
+        repo = git.Repo(path)
+        remote = repo.remotes.origin.url.replace('.git', '')  # i.e. 'https://github.com/ultralytics/yolov5'
+        commit = repo.head.commit.hexsha  # i.e. '3134699c73af83aac2a481435550b968d5792c0d'
+        try:
+            branch = repo.active_branch.name  # i.e. 'main'
+        except TypeError:  # not on any branch
+            branch = None  # i.e. 'detached HEAD' state
+        return {'remote': remote, 'branch': branch, 'commit': commit}
+    except git.exc.InvalidGitRepositoryError:  # path is not a git dir
+        return {'remote': None, 'branch': None, 'commit': None}
+
+
 def check_python(minimum='3.7.0'):
     # Check current python version vs. required python version
     check_version(platform.python_version(), minimum, name='Python ', hard=True)
@@ -1121,4 +1138,4 @@ def imshow(path, im):
 cv2.imread, cv2.imwrite, cv2.imshow = imread, imwrite, imshow  # redefine
 
 # Variables ------------------------------------------------------------------------------------------------------------
-NCOLS = 0 if is_docker() else shutil.get_terminal_size().columns  # terminal window size for tqdm
+GIT = check_git()  # repo, branch, commit