align loss to darknet
This commit is contained in:
parent
292af1f2f4
commit
750f528bfe
|
@ -137,8 +137,8 @@ class YOLOLayer(nn.Module):
|
|||
|
||||
# Training
|
||||
if targets is not None:
|
||||
MSELoss = nn.MSELoss(size_average=False)
|
||||
BCEWithLogitsLoss = nn.BCEWithLogitsLoss(size_average=False)
|
||||
MSELoss = nn.MSELoss(size_average=True)
|
||||
BCEWithLogitsLoss = nn.BCEWithLogitsLoss(size_average=True)
|
||||
# CrossEntropyLoss = nn.CrossEntropyLoss()
|
||||
|
||||
if requestPrecision:
|
||||
|
|
18
train.py
18
train.py
|
@ -6,12 +6,12 @@ from utils.datasets import *
|
|||
from utils.utils import *
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-epochs', type=int, default=160, help='number of epochs')
|
||||
parser.add_argument('-epochs', type=int, default=1, help='number of epochs')
|
||||
parser.add_argument('-batch_size', type=int, default=12, help='size of each image batch')
|
||||
parser.add_argument('-data_config_path', type=str, default='cfg/coco.data', help='data config file path')
|
||||
parser.add_argument('-cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
||||
parser.add_argument('-img_size', type=int, default=32 * 13, help='size of each image dimension')
|
||||
parser.add_argument('-resume', default=False, help='resume training flag')
|
||||
parser.add_argument('-resume', default=True, help='resume training flag')
|
||||
opt = parser.parse_args()
|
||||
print(opt)
|
||||
|
||||
|
@ -48,7 +48,7 @@ def main(opt):
|
|||
start_epoch = 0
|
||||
best_loss = float('inf')
|
||||
if opt.resume:
|
||||
checkpoint = torch.load('checkpoints/latest.pt', map_location='cpu')
|
||||
checkpoint = torch.load('checkpoints/yolov3.pt', map_location='cpu')
|
||||
|
||||
model.load_state_dict(checkpoint['model'])
|
||||
if torch.cuda.device_count() > 1:
|
||||
|
@ -115,12 +115,12 @@ def main(opt):
|
|||
continue
|
||||
|
||||
# SGD burn-in
|
||||
if (epoch == 0) & (i <= 1000):
|
||||
power = 4
|
||||
lr = 1e-3 * (i / 1000) ** power
|
||||
for g in optimizer.param_groups:
|
||||
g['lr'] = lr
|
||||
# print('SGD Burn-In LR = %9.5g' % lr, end='')
|
||||
# if (epoch == 0) & (i <= 1000):
|
||||
# power = 4
|
||||
# lr = 1e-3 * (i / 1000) ** power
|
||||
# for g in optimizer.param_groups:
|
||||
# g['lr'] = lr
|
||||
# # print('SGD Burn-In LR = %9.5g' % lr, end='')
|
||||
|
||||
# Compute loss, compute gradient, update parameters
|
||||
loss = model(imgs.to(device), targets, requestPrecision=True)
|
||||
|
|
Loading…
Reference in New Issue