multi gpus ?
qxin opened this issue · comments
qxin commented
Hi, thx ~for the beautiful code, how to use multi GPU ?
Yu Hao commented
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)
qxin commented
ok,ok, thx~~