This commit is contained in:
Glenn Jocher 2019-07-08 19:26:46 +02:00
parent 0bd763f528
commit bb1e551150
2 changed files with 11 additions and 4 deletions

View File

@ -354,7 +354,13 @@ def convert(cfg='cfg/yolov3-spp.cfg', weights='weights/yolov3-spp.weights'):
elif weights.endswith('.weights'): # darknet format
_ = load_darknet_weights(model, weights)
chkpt = {'epoch': -1, 'best_loss': None, 'model': model.state_dict(), 'optimizer': None}
chkpt = {'epoch': -1,
'best_fitness': None,
'training_results': None,
'model': model.state_dict(),
'optimizer': None}
torch.save(chkpt, 'converted.pt')
print("Success: converted '%s' to 'converted.pt'" % weights)

View File

@ -81,8 +81,9 @@ def train(
optimizer.load_state_dict(chkpt['optimizer'])
best_fitness = chkpt['best_fitness']
with open('results.txt', 'w') as file:
file.write(chkpt['training_results']) # write results.txt
if chkpt['training_results'] is not None:
with open('results.txt', 'w') as file:
file.write(chkpt['training_results']) # write results.txt
start_epoch = chkpt['epoch'] + 1
del chkpt
@ -251,7 +252,7 @@ def train(
best_fitness = fitness
# Save training results
save = (not opt.nosave) or (epoch == epochs - 1)
save = (not opt.nosave) or ((not opt.evolve) and (epoch == epochs - 1))
if save:
with open('results.txt', 'r') as file:
# Create checkpoint