updates
This commit is contained in:
parent
12e605165e
commit
ac22a717f1
7
train.py
7
train.py
|
@ -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')
|
||||||
|
|
Loading…
Reference in New Issue