Skip to content

Commit

Permalink
updates
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Nov 27, 2018
1 parent ab9ee6a commit b07ee41
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 20 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ The https://github.com/ultralytics/yolov3 repo contains inference and training c

# Requirements

Python 3.6 or later with the following `pip3 install -U -r requirements.txt` packages:
Python 3.7 or later with the following `pip3 install -U -r requirements.txt` packages:

- `numpy`
- `torch`
Expand Down
18 changes: 15 additions & 3 deletions models.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,8 +102,9 @@ def __init__(self, anchors, nC, img_dim, anchor_idxs):
self.weights = class_weights()

self.loss_means = torch.ones(6)
self.tx, self.ty, self.tw, self.th = [], [], [], []

def forward(self, p, targets=None, batch_report=False):
def forward(self, p, targets=None, batch_report=False, var=None):
FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor

bs = p.shape[0] # batch size
Expand Down Expand Up @@ -171,6 +172,17 @@ def forward(self, p, targets=None, batch_report=False):
lw = k * MSELoss(w[mask], tw[mask])
lh = k * MSELoss(h[mask], th[mask])

# self.tx.extend(tx[mask].data.numpy())
# self.ty.extend(ty[mask].data.numpy())
# self.tw.extend(tw[mask].data.numpy())
# self.th.extend(th[mask].data.numpy())
# print([np.mean(self.tx), np.std(self.tx)],[np.mean(self.ty), np.std(self.ty)],[np.mean(self.tw), np.std(self.tw)],[np.mean(self.th), np.std(self.th)])
# [0.5040668, 0.2885492] [0.51384246, 0.28328574] [-0.4754091, 0.57951087] [-0.25998235, 0.44858757]
# [0.50184494, 0.2858976] [0.51747805, 0.2896323] [0.12962963, 0.6263085] [-0.2722081, 0.61574113]
# [0.5032071, 0.28825334] [0.5063132, 0.2808862] [0.21124361, 0.44760725] [0.35445485, 0.6427766]
# import matplotlib.pyplot as plt
# plt.hist(self.x)

# lconf = k * BCEWithLogitsLoss(pred_conf[mask], mask[mask].float())
lconf = (k * 10) * BCEWithLogitsLoss(pred_conf, mask.float())

Expand Down Expand Up @@ -227,7 +239,7 @@ def __init__(self, cfg_path, img_size=416):
self.img_size = img_size
self.loss_names = ['loss', 'x', 'y', 'w', 'h', 'conf', 'cls', 'nT', 'TP', 'FP', 'FPe', 'FN', 'TC']

def forward(self, x, targets=None, batch_report=False):
def forward(self, x, targets=None, batch_report=False, var=0):
is_training = targets is not None
output = []
self.losses = defaultdict(float)
Expand All @@ -245,7 +257,7 @@ def forward(self, x, targets=None, batch_report=False):
elif module_def['type'] == 'yolo':
# Train phase: get loss
if is_training:
x, *losses = module[0](x, targets, batch_report)
x, *losses = module[0](x, targets, batch_report, var)
for name, loss in zip(self.loss_names, losses):
self.losses[name] += loss
# Test phase: Get detections
Expand Down
40 changes: 24 additions & 16 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@
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('-batch_report', default=False, help='report TP, FP, FN, P and R per batch (slower)')
parser.add_argument('-optimizer', default='SGD', help='Optimizer')
parser.add_argument('-optimizer', default='SGD', help='optimizer')
parser.add_argument('-freeze_darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
parser.add_argument('-var', type=float, default=0, help='optional test variable')
opt = parser.parse_args()
print(opt)

Expand Down Expand Up @@ -51,9 +53,7 @@ def main(opt):
# Get dataloader
dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=opt.img_size, augment=True)

# Reload saved optimizer state
start_epoch = 0
best_loss = float('inf')
lr0 = 0.001
if opt.resume:
checkpoint = torch.load('weights/latest.pt', map_location='cpu')

Expand All @@ -69,10 +69,7 @@ def main(opt):
# p.requires_grad = False

# Set optimizer
if opt.optimizer is 'Adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4, weight_decay=5e-4)
else:
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3, momentum=.9, weight_decay=5e-4)
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=lr0, momentum=.9)

start_epoch = checkpoint['epoch'] + 1
if checkpoint['optimizer'] is not None:
Expand All @@ -82,6 +79,9 @@ def main(opt):
del checkpoint # current, saved

else:
start_epoch = 0
best_loss = float('inf')

# Initialize model with darknet53 weights (optional)
if not os.path.isfile('weights/darknet53.conv.74'):
os.system('wget https://pjreddie.com/media/files/darknet53.conv.74 -P weights')
Expand All @@ -93,10 +93,7 @@ def main(opt):
model.to(device).train()

# Set optimizer
if opt.optimizer is 'Adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4, weight_decay=5e-4)
else:
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3, momentum=.9, weight_decay=5e-4)
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=lr0, momentum=.9)

# Set scheduler
# scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[54, 61], gamma=0.1)
Expand All @@ -114,12 +111,23 @@ def main(opt):

# Update scheduler (manual) at 0, 54, 61 epochs to 1e-3, 1e-4, 1e-5
if epoch > 50:
lr = 1e-4
lr = lr0 / 10
else:
lr = 1e-3
lr = lr0
for g in optimizer.param_groups:
g['lr'] = lr

# Freeze darknet53.conv.74 layers for first epoch
if opt.freeze_darknet53:
if epoch == 0:
for i, (name, p) in enumerate(model.named_parameters()):
if int(name.split('.')[1]) < 75: # if layer < 75
p.requires_grad = False
elif epoch == 1:
for i, (name, p) in enumerate(model.named_parameters()):
if int(name.split('.')[1]) < 75: # if layer < 75
p.requires_grad = True

ui = -1
rloss = defaultdict(float) # running loss
metrics = torch.zeros(3, num_classes)
Expand All @@ -130,12 +138,12 @@ def main(opt):

# SGD burn-in
if (epoch == 0) & (i <= 1000):
lr = 1e-3 * (i / 1000) ** 4
lr = lr0 * (i / 1000) ** 4
for g in optimizer.param_groups:
g['lr'] = lr

# Compute loss, compute gradient, update parameters
loss = model(imgs.to(device), targets, batch_report=opt.batch_report)
loss = model(imgs.to(device), targets, batch_report=opt.batch_report, var=opt.var)
loss.backward()

# accumulated_batches = 1 # accumulate gradient for 4 batches before stepping optimizer
Expand Down

0 comments on commit b07ee41

Please sign in to comment.