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

updates

This commit is contained in:
Glenn Jocher 2018-11-16 20:01:38 +01:00
parent a021f97110
commit d2c5d7a5fd
1 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
models.py
View File

@ -101,6 +101,9 @@ class YOLOLayer(nn.Module):
self.anchor_h = self.scaled_anchors[:, 1:2].view((1, nA, 1, 1)) self.anchor_h = self.scaled_anchors[:, 1:2].view((1, nA, 1, 1))
self.weights = class_weights() self.weights = class_weights()
self.batch_count = 0
self.loss_means = torch.zeros(6)
def forward(self, p, targets=None, requestPrecision=False): def forward(self, p, targets=None, requestPrecision=False):
FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor
@ -139,6 +142,7 @@ class YOLOLayer(nn.Module):
# Training # Training
if targets is not None: if targets is not None:
self.batch_count += 1
MSELoss = nn.MSELoss() MSELoss = nn.MSELoss()
BCEWithLogitsLoss = nn.BCEWithLogitsLoss() BCEWithLogitsLoss = nn.BCEWithLogitsLoss()
CrossEntropyLoss = nn.CrossEntropyLoss() CrossEntropyLoss = nn.CrossEntropyLoss()
@ -181,6 +185,14 @@ class YOLOLayer(nn.Module):
# lconf += k * BCEWithLogitsLoss(pred_conf[~mask], mask[~mask].float()) # lconf += k * BCEWithLogitsLoss(pred_conf[~mask], mask[~mask].float())
# Sum loss components # Sum loss components
balance_losses_flag = True
if balance_losses_flag:
loss_vec = torch.FloatTensor([lx.data, ly.data, lw.data, lh.data, lconf.data, lcls.data])
self.loss_means = self.loss_means * 0.99 + loss_vec * 0.01
k = 1 / self.loss_means.clone()
k /= k.sum()
loss = (lx * k[0] + ly * k[1] + lw * k[2] + lh * k[3] + lconf * k[4] + lcls * k[5]) * loss_vec.sum()
else:
loss = lx + ly + lw + lh + lconf + lcls loss = lx + ly + lw + lh + lconf + lcls
# Sum False Positives from unassigned anchors # Sum False Positives from unassigned anchors