updates
This commit is contained in:
parent
0bd763f528
commit
bb1e551150
|
@ -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)
|
||||
|
||||
|
|
7
train.py
7
train.py
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue