From 6e19245dc8dd9a16d8e48a9b9493f53384b8bbd1 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 9 Apr 2020 19:53:29 -0700 Subject: [PATCH] auto strip optimizer from best.pt after training --- train.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/train.py b/train.py index 5c470e3e..04292736 100644 --- a/train.py +++ b/train.py @@ -362,14 +362,13 @@ def train(): n = opt.name if len(n): n = '_' + n if not n.isnumeric() else n - fresults, flast, fbest = 'results%s.txt' % n, 'last%s.pt' % n, 'best%s.pt' % n - os.rename('results.txt', fresults) - os.rename(wdir + 'last.pt', wdir + flast) if os.path.exists(wdir + 'last.pt') else None - os.rename(wdir + 'best.pt', wdir + fbest) if os.path.exists(wdir + 'best.pt') else None - if opt.bucket: # save to cloud - os.system('gsutil cp %s gs://%s/results' % (fresults, opt.bucket)) - os.system('gsutil cp %s gs://%s/weights' % (wdir + flast, opt.bucket)) - os.system('gsutil cp %s gs://%s/weights' % (wdir + fbest, opt.bucket)) + fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n + for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'], [flast, fbest, fresults]): + if os.path.exists(f1): + os.rename(f1, f2) # rename + ispt = f2.endswith('.pt') # is *.pt + strip_optimizer(f2) if ispt else None # strip optimizer + os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket and ispt else None # upload if not opt.evolve: plot_results() # save as results.png