This commit is contained in:
Glenn Jocher 2019-02-22 16:15:20 +01:00
parent 12e605165e
commit ac22a717f1
1 changed files with 4 additions and 5 deletions

View File

@ -40,8 +40,10 @@ def train(
lr0 = 0.001 lr0 = 0.001
cutoff = -1 # backbone reaches to cutoff layer cutoff = -1 # backbone reaches to cutoff layer
start_epoch = 0
best_loss = float('inf')
if resume: if resume:
checkpoint = torch.load(latest, map_location='cpu') checkpoint = torch.load('weights/yolov3.pt', map_location='cpu')
# Load weights to resume from # Load weights to resume from
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
@ -65,9 +67,6 @@ def train(
del checkpoint # current, saved del checkpoint # current, saved
else: else:
start_epoch = 0
best_loss = float('inf')
# Initialize model with backbone (optional) # Initialize model with backbone (optional)
if cfg.endswith('yolov3.cfg'): if cfg.endswith('yolov3.cfg'):
load_darknet_weights(model, weights + 'darknet53.conv.74') load_darknet_weights(model, weights + 'darknet53.conv.74')