Update center_loss.py

Loop structure is very time consuming

layers/center_loss.py

@@ -44,13 +44,15 @@ class CenterLoss(nn.Module):
         labels = labels.unsqueeze(1).expand(batch_size, self.num_classes)
         mask = labels.eq(classes.expand(batch_size, self.num_classes))
 
-        dist = []
+        dist = distmat * mask.float()
-        for i in range(batch_size):
+        loss = dist.clamp(min=1e-12, max=1e+12).sum() / batch_size
-            value = distmat[i][mask[i]]
+        #dist = []
-            value = value.clamp(min=1e-12, max=1e+12)  # for numerical stability
+        #for i in range(batch_size):
-            dist.append(value)
+        #    value = distmat[i][mask[i]]
-        dist = torch.cat(dist)
+        #    value = value.clamp(min=1e-12, max=1e+12)  # for numerical stability
-        loss = dist.mean()
+        #    dist.append(value)
+        #dist = torch.cat(dist)
+        #loss = dist.mean()
         return loss