diff options
| -rw-r--r-- | models/model.py | 2 | ||||
| -rw-r--r-- | models/rgb_part_net.py | 12 | ||||
| -rw-r--r-- | utils/triplet_loss.py | 2 |
3 files changed, 11 insertions, 5 deletions
diff --git a/models/model.py b/models/model.py index 9cb6a8e..f79b832 100644 --- a/models/model.py +++ b/models/model.py @@ -193,6 +193,8 @@ class Model: x_c1 = batch_c1['clip'].to(self.device) x_c2 = batch_c2['clip'].to(self.device) y = batch_c1['label'].to(self.device) + # Duplicate labels for each part + y = y.unsqueeze(1).repeat(1, self.rgb_pn.num_total_parts) losses, images = self.rgb_pn(x_c1, x_c2, y) loss = losses.sum() loss.backward() diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index 8ebcfd3..2aa680c 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -44,7 +44,8 @@ class RGBPartNet(nn.Module): ae_feature_channels * 2, out_channels, hpm_use_1x1conv, hpm_scales, hpm_use_avg_pool, hpm_use_max_pool ) - empty_fc = torch.empty(self.hpm_num_parts + tfa_num_parts, + self.num_total_parts = self.hpm_num_parts + tfa_num_parts + empty_fc = torch.empty(self.num_total_parts, out_channels, embedding_dims) self.fc_mat = nn.Parameter(empty_fc) @@ -75,8 +76,13 @@ class RGBPartNet(nn.Module): x = self.fc(x) if self.training: - hpm_ba_trip = self.hpm_ba_trip(x[:self.hpm_num_parts], y) - pn_ba_trip = self.pn_ba_trip(x[self.hpm_num_parts:], y) + y = y.T + hpm_ba_trip = self.hpm_ba_trip( + x[:self.hpm_num_parts], y[:self.hpm_num_parts] + ) + pn_ba_trip = self.pn_ba_trip( + x[self.hpm_num_parts:], y[self.hpm_num_parts:] + ) losses = torch.stack((*losses, hpm_ba_trip, pn_ba_trip)) return losses, images else: diff --git a/utils/triplet_loss.py b/utils/triplet_loss.py index d573ef4..954def2 100644 --- a/utils/triplet_loss.py +++ b/utils/triplet_loss.py @@ -9,9 +9,7 @@ class BatchAllTripletLoss(nn.Module): self.margin = margin def forward(self, x, y): - # Duplicate labels for each part p, n, c = x.size() - y = y.repeat(p, 1) # Euclidean distance p x n x n x_squared_sum = torch.sum(x ** 2, dim=2) |