longcw / yolo2-pytorch

YOLOv2 in PyTorch

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

multi gpus ?

qxin opened this issue · comments

commented

Hi, thx ~for the beautiful code, how to use multi GPU ?

I failed...

It seems that we should calculate loss outside the forward function. Use net.module.loss directly will cause None add None error.


Update

if torch.cuda.device_count() > 1:
    net = nn.DataParallel(net)
    @property
    def loss(self):
        if self.bbox_loss is not None and self.iou_loss is not None and self.cls_loss is not None:
            return self.bbox_loss, self.iou_loss, self.cls_loss
        return None, None, None
def forward(self, im_data, gt_boxes=None, gt_classes=None, dontcare=None, size_index=0):
    #...
    return bbox_pred, iou_pred, prob_pred, self.loss
            bbox_pred, iou_pred, prob_pred, losses = net(im_data, batch['gt_boxes'], batch['gt_classes'],
                                                         batch['dontcare'], size_index)
            bbox_loss = torch.sum(losses[0])
            iou_loss = torch.sum(losses[1])
            cls_loss = torch.sum(losses[2])
            # Backward
            optimizer.zero_grad()
            loss = bbox_loss + iou_loss + cls_loss
            loss.backward()
            optimizer.step()

But I can not make the images during training.

No speed boost unless you use a REAL Dataset (not my version)

commented

ok,ok, thx~~