longcw / yolo2-pytorch

YOLOv2 in PyTorch
1.54k stars 421 forks source link

Training issue #88

Open leonardoaraujosantos opened 6 years ago

leonardoaraujosantos commented 6 years ago

Hi, I've cloned the code today and try the training and got the following error

python3 train.py
/home/leoara01/work/yolo2-pytorch/darknet.py:214: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
  prob_pred = F.softmax(score_pred.view(-1, score_pred.size()[-1])).view_as(score_pred)  # noqa
Traceback (most recent call last):
  File "train.py", line 78, in <module>
    bbox_loss += net.bbox_loss.data.cpu().numpy()[0]
IndexError: too many indices for array

I'm using pytorch 0.4

longcw commented 6 years ago

https://github.com/longcw/yolo2-pytorch/issues/59

leonardoaraujosantos commented 6 years ago

Thanks, I will check that now

leonardoaraujosantos commented 6 years ago

Hi @longcw now it works, here are the changes

leoara01@linuxdev:~/work/yolo2-pytorch$ git diff train.py
diff --git a/train.py b/train.py
index 8c33cb4..51ba809 100644
--- a/train.py
+++ b/train.py
@@ -75,10 +75,10 @@ for step in range(start_epoch * imdb.batch_per_epoch,

     # backward
     loss = net.loss
-    bbox_loss += net.bbox_loss.data.cpu().numpy()[0]
-    iou_loss += net.iou_loss.data.cpu().numpy()[0]
-    cls_loss += net.cls_loss.data.cpu().numpy()[0]
-    train_loss += loss.data.cpu().numpy()[0]
+    bbox_loss += net.bbox_loss.data.cpu().item()
+    iou_loss += net.iou_loss.data.cpu().item()
+    cls_loss += net.cls_loss.data.cpu().item()
+    train_loss += loss.data.cpu().item()
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()

Would you like that I create a Pull request for the issue?