auto strip optimizer from best.pt after training
This commit is contained in:
		
							parent
							
								
									bc74822540
								
							
						
					
					
						commit
						6e19245dc8
					
				
							
								
								
									
										15
									
								
								train.py
								
								
								
								
							
							
						
						
									
										15
									
								
								train.py
								
								
								
								
							| 
						 | 
					@ -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
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
		Loading…
	
		Reference in New Issue