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 | |
| parent | 7825f978f198e56958703f0d08f7ccbd8cef49ca (diff) | |
| parent | 36cf502afe9b93efe31c244030270b0a62e644b8 (diff) | |
Merge branch 'master' into python3.8
# Conflicts:
# models/model.py
Diffstat (limited to 'utils')
| -rw-r--r-- | utils/misc.py | 10 | ||||
| -rw-r--r-- | utils/triplet_loss.py | 9 |
2 files changed, 15 insertions, 4 deletions
diff --git a/utils/misc.py b/utils/misc.py new file mode 100644 index 0000000..b850830 --- /dev/null +++ b/utils/misc.py @@ -0,0 +1,10 @@ +import os + +from utils.configuration import SystemConfiguration + + +def set_visible_cuda(config: SystemConfiguration): + """Set environment variable CUDA device(s)""" + CUDA_VISIBLE_DEVICES = config.get('CUDA_VISIBLE_DEVICES', None) + if CUDA_VISIBLE_DEVICES: + os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_VISIBLE_DEVICES 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 |