glenn-jocher commited on
Commit
3bdea3f
·
1 Parent(s): 6ca3f35

strip_optimizer() bug fix #253

Browse files
Files changed (1) hide show
  1. train.py +11 -11
train.py CHANGED
@@ -336,17 +336,17 @@ def train(hyp):
336
  # end epoch ----------------------------------------------------------------------------------------------------
337
  # end training
338
 
339
- n = opt.name
340
- if len(n):
341
- n = '_' + n if not n.isnumeric() else n
342
- fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n
343
- for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'], [flast, fbest, fresults]):
344
- if os.path.exists(f1):
345
- os.rename(f1, f2) # rename
346
- ispt = f2.endswith('.pt') # is *.pt
347
- strip_optimizer(f2) if ispt else None # strip optimizer
348
- os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket and ispt else None # upload
349
-
350
  if not opt.evolve:
351
  plot_results() # save as results.png
352
  print('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
 
336
  # end epoch ----------------------------------------------------------------------------------------------------
337
  # end training
338
 
339
+ # Strip optimizers
340
+ n = ('_' if len(opt.name) and not opt.name.isnumeric() else '') + opt.name
341
+ fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n
342
+ for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'], [flast, fbest, fresults]):
343
+ if os.path.exists(f1):
344
+ os.rename(f1, f2) # rename
345
+ ispt = f2.endswith('.pt') # is *.pt
346
+ strip_optimizer(f2) if ispt else None # strip optimizer
347
+ os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket and ispt else None # upload
348
+
349
+ # Finish
350
  if not opt.evolve:
351
  plot_results() # save as results.png
352
  print('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))