jasder
/
yolov5
1
0
Fork 0
Code Issues Pull Requests 1 Releases Wiki Activity

Update train.py

Browse Source
pull/1/head
Glenn Jocher 5 years ago committed by GitHub
parent 22ab1c295f
commit 6b134d93c5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 20 additions and 39 deletions
Show Stats Download Patch File Download Diff File

51
train.py
Unescape View File

@ -44,11 +44,8 @@ hyp = {'optimizer': 'SGD', # ['adam', 'SGD', None] if none, default is SGD
def train(hyp):
#write all results to the tb log_dir, so all data from one run is together
log_dir = tb_writer.log_dir
#weights dir unique to each experiment
wdir = os.path.join(log_dir, 'weights') + os.sep # weights dir
log_dir = tb_writer.log_dir # run directory
wdir = str(Path(log_dir) / 'weights') + os.sep # weights directory
os.makedirs(wdir, exist_ok=True)
last = wdir + 'last.pt'
@ -92,8 +89,8 @@ def train(hyp):
else:
pg0.append(v) # all else
if hyp['optimizer'] =='adam':
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) #use default beta2, adjust beta1 for Adam momentum per momentum adjustments in https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
if hyp['optimizer'] == 'adam': # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
@ -177,11 +174,10 @@ def train(hyp):
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
model.names = data_dict['names']
#save hyperparamter and training options in run folder
with open(os.path.join(log_dir, 'hyp.yaml'), 'w') as f:
# Save run settings
with open(Path(log_dir) / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(os.path.join(log_dir, 'opt.yaml'), 'w') as f:
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Class frequency
@ -189,14 +185,10 @@ def train(hyp):
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1.
# model._initialize_biases(cf.to(device))
#always plot labels to log_dir
plot_labels(labels, save_dir=log_dir)
if tb_writer:
tb_writer.add_histogram('classes', c, 0)
# Check anchors
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
@ -284,7 +276,7 @@ def train(hyp):
# Plot
if ni < 3:
f = os.path.join(log_dir, 'train_batch%g.jpg' % ni) # filename
f = str(Path(log_dir) / ('train_batch%g.jpg' % ni)) # filename
result = plot_images(images=imgs, targets=targets, paths=paths, fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
@ -368,14 +360,14 @@ def train(hyp):
if __name__ == '__main__':
check_git_status()
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model cfg path[*.yaml]')
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data cfg path [*.yaml]')
parser.add_argument('--hyp', type=str, default='',help='hyp cfg path [*.yaml].')
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
parser.add_argument('--hyp', type=str, default='', help='hyp.yaml path (optional)')
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes')
parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--resume', nargs='?', const = 'get_last', default=False, help='resume training from given path/to/last.pt, or most recent run if blank.')
parser.add_argument('--resume', nargs='?', const = 'get_last', default=False, help='resume from given path/to/last.pt, or most recent run if blank.')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--notest', action='store_true', help='only test final epoch')
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
@ -387,20 +379,15 @@ if __name__ == '__main__':
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
opt = parser.parse_args()
# use given path/to/last.pt or find most recent run if no path given
last = get_latest_run() if opt.resume == 'get_last' else opt.resume
last = get_latest_run() if opt.resume == 'get_last' else opt.resume # resume from most recent run
if last and not opt.weights:
print(f'Resuming training from {last}')
opt.weights = last if opt.resume and not opt.weights else opt.weights
opt.cfg = check_file(opt.cfg) # check file
opt.data = check_file(opt.data) # check file
opt.hyp = check_file(opt.hyp) if opt.hyp else '' # check file
print(opt)
opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test)
device = torch_utils.select_device(opt.device, apex=mixed_precision, batch_size=opt.batch_size)
@ -410,16 +397,10 @@ if __name__ == '__main__':
# Train
if not opt.evolve:
tb_writer = SummaryWriter(comment=opt.name)
#updates hyp defaults from hyp.yaml
if opt.hyp:
if opt.hyp: # update hyps
with open(opt.hyp) as f:
updated_hyp = yaml.load(f, Loader=yaml.FullLoader)
hyp.update(updated_hyp)
hyp.update(yaml.load(f, Loader=yaml.FullLoader))
# Print focal loss if gamma > 0
if hyp['fl_gamma']:
print('Using FocalLoss(gamma=%g)' % hyp['fl_gamma'])
print(f'Beginning training with {hyp}\n\n')
print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')

Write Preview
Loading…
Cancel
Save