centerloss.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. from __future__ import absolute_import
  2. from __future__ import division
  3. from __future__ import print_function
  4. import paddle
  5. import paddle.nn as nn
  6. import paddle.nn.functional as F
  7. class CenterLoss(nn.Layer):
  8. def __init__(self, num_classes=5013, feat_dim=2048):
  9. super(CenterLoss, self).__init__()
  10. self.num_classes = num_classes
  11. self.feat_dim = feat_dim
  12. self.centers = paddle.randn(
  13. shape=[self.num_classes, self.feat_dim]).astype(
  14. "float64") #random center
  15. def __call__(self, input, target):
  16. """
  17. inputs: network output: {"features: xxx", "logits": xxxx}
  18. target: image label
  19. """
  20. feats = input["features"]
  21. labels = target
  22. batch_size = feats.shape[0]
  23. #calc feat * feat
  24. dist1 = paddle.sum(paddle.square(feats), axis=1, keepdim=True)
  25. dist1 = paddle.expand(dist1, [batch_size, self.num_classes])
  26. #dist2 of centers
  27. dist2 = paddle.sum(paddle.square(self.centers), axis=1,
  28. keepdim=True) #num_classes
  29. dist2 = paddle.expand(dist2,
  30. [self.num_classes, batch_size]).astype("float64")
  31. dist2 = paddle.transpose(dist2, [1, 0])
  32. #first x * x + y * y
  33. distmat = paddle.add(dist1, dist2)
  34. tmp = paddle.matmul(feats, paddle.transpose(self.centers, [1, 0]))
  35. distmat = distmat - 2.0 * tmp
  36. #generate the mask
  37. classes = paddle.arange(self.num_classes).astype("int64")
  38. labels = paddle.expand(
  39. paddle.unsqueeze(labels, 1), (batch_size, self.num_classes))
  40. mask = paddle.equal(
  41. paddle.expand(classes, [batch_size, self.num_classes]),
  42. labels).astype("float64") #get mask
  43. dist = paddle.multiply(distmat, mask)
  44. loss = paddle.sum(paddle.clip(dist, min=1e-12, max=1e+12)) / batch_size
  45. return {'CenterLoss': loss}