diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-01-12 11:29:02 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-01-12 11:29:02 +0800 |
| commit | 966d4431c037b0c4641aa2a5fc22f05be064b331 (patch) | |
| tree | 0239ba89d31857a7f086acf627fc1bbf167855a9 /utils/triplet_loss.py | |
| parent | 7825f978f198e56958703f0d08f7ccbd8cef49ca (diff) | |
| parent | 36cf502afe9b93efe31c244030270b0a62e644b8 (diff) | |
Merge branch 'master' into python3.8
# Conflicts:
# models/model.py
Diffstat (limited to 'utils/triplet_loss.py')
| -rw-r--r-- | utils/triplet_loss.py | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/utils/triplet_loss.py b/utils/triplet_loss.py index 1d63a0e..8c143d6 100644 --- a/utils/triplet_loss.py +++ b/utils/triplet_loss.py @@ -15,8 +15,8 @@ class BatchAllTripletLoss(nn.Module): # Euclidean distance p x n x n x_squared_sum = torch.sum(x ** 2, dim=2) - x1_squared_sum = x_squared_sum.unsqueeze(1) - x2_squared_sum = x_squared_sum.unsqueeze(2) + x1_squared_sum = x_squared_sum.unsqueeze(2) + x2_squared_sum = x_squared_sum.unsqueeze(1) x1_times_x2_sum = x @ x.transpose(1, 2) dist = torch.sqrt( F.relu(x1_squared_sum - 2 * x1_times_x2_sum + x2_squared_sum) @@ -30,8 +30,9 @@ class BatchAllTripletLoss(nn.Module): all_loss = F.relu(self.margin + positive_negative_dist).view(p, -1) # Non-zero parted mean - parted_loss_mean = all_loss.sum(1) / (all_loss != 0).sum(1) - parted_loss_mean[parted_loss_mean == float('Inf')] = 0 + non_zero_counts = (all_loss != 0).sum(1) + parted_loss_mean = all_loss.sum(1) / non_zero_counts + parted_loss_mean[non_zero_counts == 0] = 0 loss = parted_loss_mean.sum() return loss |