auto strip optimizer from best.pt after training

This commit is contained in:
Glenn Jocher 2020-04-09 19:53:29 -07:00
parent bc74822540
commit 6e19245dc8
1 changed files with 7 additions and 8 deletions

View File

@ -362,14 +362,13 @@ def train():
n = opt.name n = opt.name
if len(n): if len(n):
n = '_' + n if not n.isnumeric() else n n = '_' + n if not n.isnumeric() else n
fresults, flast, fbest = 'results%s.txt' % n, 'last%s.pt' % n, 'best%s.pt' % n fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n
os.rename('results.txt', fresults) for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'], [flast, fbest, fresults]):
os.rename(wdir + 'last.pt', wdir + flast) if os.path.exists(wdir + 'last.pt') else None if os.path.exists(f1):
os.rename(wdir + 'best.pt', wdir + fbest) if os.path.exists(wdir + 'best.pt') else None os.rename(f1, f2) # rename
if opt.bucket: # save to cloud ispt = f2.endswith('.pt') # is *.pt
os.system('gsutil cp %s gs://%s/results' % (fresults, opt.bucket)) strip_optimizer(f2) if ispt else None # strip optimizer
os.system('gsutil cp %s gs://%s/weights' % (wdir + flast, opt.bucket)) os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket and ispt else None # upload
os.system('gsutil cp %s gs://%s/weights' % (wdir + fbest, opt.bucket))
if not opt.evolve: if not opt.evolve:
plot_results() # save as results.png plot_results() # save as results.png