diff --git a/models.py b/models.py index 52c8cb37..c68aaad7 100755 --- a/models.py +++ b/models.py @@ -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) diff --git a/train.py b/train.py index bf6d95b2..c8aeac52 100644 --- a/train.py +++ b/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