Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-04-02 14:29:35 +02:00
parent 748ff9b5b9
commit 47400aa066
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
train.py
View File

@ -53,7 +53,7 @@ def train(
if checkpoint['optimizer'] is not None: if checkpoint['optimizer'] is not None:
optimizer.load_state_dict(checkpoint['optimizer']) optimizer.load_state_dict(checkpoint['optimizer'])
best_loss = checkpoint['best_loss'] best_loss = checkpoint['best_loss']
del checkpoint # current, saved del checkpoint
else: # Initialize model with backbone (optional) else: # Initialize model with backbone (optional)
if cfg.endswith('yolov3.cfg'): if cfg.endswith('yolov3.cfg'):
@ -180,6 +180,7 @@ def train(
model) is nn.parallel.DistributedDataParallel else model.state_dict(), model) is nn.parallel.DistributedDataParallel else model.state_dict(),
'optimizer': optimizer.state_dict()} 'optimizer': optimizer.state_dict()}
torch.save(checkpoint, latest) torch.save(checkpoint, latest)
del checkpoint
# Save best checkpoint # Save best checkpoint
if best_loss == mloss['total']: if best_loss == mloss['total']: