BCE to CE lconf + batch size 16

This commit is contained in:
Glenn Jocher 2018-10-11 17:43:34 +02:00
parent d336e0053d
commit 24a41972cb
2 changed files with 3 additions and 5 deletions

View File

@ -101,7 +101,6 @@ class YOLOLayer(nn.Module):
self.anchor_h = self.scaled_anchors[:, 1:2].view((1, nA, 1, 1))
self.weights = class_weights()
def forward(self, p, targets=None, requestPrecision=False):
FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor
@ -173,8 +172,8 @@ class YOLOLayer(nn.Module):
# lconf = k * BCEWithLogitsLoss(pred_conf[mask], mask[mask].float())
lconf = k * BCEWithLogitsLoss(pred_conf, mask.float())
# lcls = k * CrossEntropyLoss(pred_cls[mask], torch.argmax(tcls, 1))
lcls = k * BCEWithLogitsLoss(pred_cls[mask], tcls.float())
lcls = k * CrossEntropyLoss(pred_cls[mask], torch.argmax(tcls, 1))
# lcls = k * BCEWithLogitsLoss(pred_cls[mask], tcls.float())
else:
lx, ly, lw, lh, lcls, lconf = FT([0]), FT([0]), FT([0]), FT([0]), FT([0]), FT([0])

View File

@ -26,7 +26,6 @@ cd yolov3/checkpoints
wget https://storage.googleapis.com/ultralytics/yolov3.pt
cp yolov3.pt latest.pt
cd ..
python3 train.py -img_size 416 -batch_size 12 -epochs 1 -resume 1
python3 train.py -img_size 416 -batch_size 16 -epochs 1 -resume 1
python3 test.py -img_size 416 -weights_path checkpoints/latest.pt -conf_thres 0.5