From 820d3dec284f38e6a3089dad5277bc3f6c5123bf Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sat, 20 Feb 2021 14:19:30 +0800 Subject: Separate triplet loss from model --- models/rgb_part_net.py | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) (limited to 'models/rgb_part_net.py') diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index 67acac3..408bca0 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -4,7 +4,6 @@ import torch.nn as nn from models.auto_encoder import AutoEncoder from models.hpm import HorizontalPyramidMatching from models.part_net import PartNet -from utils.triplet_loss import BatchAllTripletLoss class RGBPartNet(nn.Module): @@ -25,7 +24,6 @@ class RGBPartNet(nn.Module): tfa_squeeze_ratio: int = 4, tfa_num_parts: int = 16, embedding_dims: int = 256, - triplet_margins: tuple[float, float] = (0.2, 0.2), image_log_on: bool = False ): super().__init__() @@ -50,17 +48,13 @@ class RGBPartNet(nn.Module): out_channels, embedding_dims) self.fc_mat = nn.Parameter(empty_fc) - (hpm_margin, pn_margin) = triplet_margins - self.hpm_ba_trip = BatchAllTripletLoss(hpm_margin) - self.pn_ba_trip = BatchAllTripletLoss(pn_margin) - def fc(self, x): return x @ self.fc_mat - def forward(self, x_c1, x_c2=None, y=None): + def forward(self, x_c1, x_c2=None): # Step 1: Disentanglement # n, t, c, h, w - ((x_c, x_p), losses, images) = self._disentangle(x_c1, x_c2) + ((x_c, x_p), ae_losses, images) = self._disentangle(x_c1, x_c2) # Step 2.a: Static Gait Feature Aggregation & HPM # n, c, h, w @@ -77,15 +71,7 @@ class RGBPartNet(nn.Module): x = self.fc(x) if self.training: - 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 + return x, ae_losses, images else: return x.unsqueeze(1).view(-1) -- cgit v1.2.3 From c52fdc2748e272a5195303299a9739291be32281 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 21 Feb 2021 19:00:30 +0800 Subject: Remove FConv blocks --- models/rgb_part_net.py | 37 ++++++++++++++++++------------------- 1 file changed, 18 insertions(+), 19 deletions(-) (limited to 'models/rgb_part_net.py') diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index 408bca0..936ec46 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -17,16 +17,13 @@ class RGBPartNet(nn.Module): hpm_scales: tuple[int, ...] = (1, 2, 4), hpm_use_avg_pool: bool = True, hpm_use_max_pool: bool = True, - fpfe_feature_channels: int = 32, - fpfe_kernel_sizes: tuple[tuple, ...] = ((5, 3), (3, 3), (3, 3)), - fpfe_paddings: tuple[tuple, ...] = ((2, 1), (1, 1), (1, 1)), - fpfe_halving: tuple[int, ...] = (0, 2, 3), tfa_squeeze_ratio: int = 4, tfa_num_parts: int = 16, embedding_dims: int = 256, image_log_on: bool = False ): super().__init__() + self.h, self.w = ae_in_size (self.f_a_dim, self.f_c_dim, self.f_p_dim) = f_a_c_p_dims self.hpm_num_parts = sum(hpm_scales) self.image_log_on = image_log_on @@ -34,18 +31,17 @@ class RGBPartNet(nn.Module): self.ae = AutoEncoder( ae_in_channels, ae_in_size, ae_feature_channels, f_a_c_p_dims ) + self.pn_in_channels = ae_feature_channels * 2 self.pn = PartNet( - ae_in_channels, fpfe_feature_channels, fpfe_kernel_sizes, - fpfe_paddings, fpfe_halving, tfa_squeeze_ratio, tfa_num_parts + self.pn_in_channels, tfa_squeeze_ratio, tfa_num_parts ) - out_channels = self.pn.tfa_in_channels self.hpm = HorizontalPyramidMatching( - ae_feature_channels * 2, out_channels, hpm_use_1x1conv, + ae_feature_channels * 2, self.pn_in_channels, hpm_use_1x1conv, hpm_scales, hpm_use_avg_pool, hpm_use_max_pool ) self.num_total_parts = self.hpm_num_parts + tfa_num_parts empty_fc = torch.empty(self.num_total_parts, - out_channels, embedding_dims) + self.pn_in_channels, embedding_dims) self.fc_mat = nn.Parameter(empty_fc) def fc(self, x): @@ -82,17 +78,20 @@ class RGBPartNet(nn.Module): if self.training: ((f_a_, f_c_, f_p_), losses) = self.ae(x_c1_t2, x_c1_t1, x_c2_t2) # Decode features - with torch.no_grad(): - x_c = self._decode_cano_feature(f_c_, n, t, device) - x_p = self._decode_pose_feature(f_p_, n, t, c, h, w, device) + x_c = self._decode_cano_feature(f_c_, n, t, device) + x_p_ = self._decode_pose_feature(f_p_, n, t, c, h, w, device) + x_p = x_p_.view(n, t, self.pn_in_channels, self.h // 4, self.w // 4) - i_a, i_c, i_p = None, None, None - if self.image_log_on: + i_a, i_c, i_p = None, None, None + if self.image_log_on: + with torch.no_grad(): i_a = self._decode_appr_feature(f_a_, n, t, device) # Continue decoding canonical features i_c = self.ae.decoder.trans_conv3(x_c) i_c = torch.sigmoid(self.ae.decoder.trans_conv4(i_c)) - i_p = x_p + i_p_ = self.ae.decoder.trans_conv3(x_p_) + i_p_ = torch.sigmoid(self.ae.decoder.trans_conv4(i_p_)) + i_p = i_p_.view(n, t, c, h, w) return (x_c, x_p), losses, (i_a, i_c, i_p) @@ -119,7 +118,7 @@ class RGBPartNet(nn.Module): torch.zeros((n, self.f_a_dim), device=device), f_c.mean(1), torch.zeros((n, self.f_p_dim), device=device), - cano_only=True + is_feature_map=True ) return x_c @@ -128,7 +127,7 @@ class RGBPartNet(nn.Module): x_p_ = self.ae.decoder( torch.zeros((n * t, self.f_a_dim), device=device), torch.zeros((n * t, self.f_c_dim), device=device), - f_p_ + f_p_, + is_feature_map=True ) - x_p = x_p_.view(n, t, c, h, w) - return x_p + return x_p_ -- cgit v1.2.3 From 390bac976ff52fe0c3cf6bea820c22084613ee94 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Fri, 26 Feb 2021 20:09:22 +0800 Subject: Fix predict function --- models/rgb_part_net.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'models/rgb_part_net.py') diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index 936ec46..4367c62 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -74,8 +74,8 @@ class RGBPartNet(nn.Module): def _disentangle(self, x_c1_t2, x_c2_t2=None): n, t, c, h, w = x_c1_t2.size() device = x_c1_t2.device - x_c1_t1 = x_c1_t2[:, torch.randperm(t), :, :, :] if self.training: + x_c1_t1 = x_c1_t2[:, torch.randperm(t), :, :, :] ((f_a_, f_c_, f_p_), losses) = self.ae(x_c1_t2, x_c1_t1, x_c2_t2) # Decode features x_c = self._decode_cano_feature(f_c_, n, t, device) @@ -98,7 +98,8 @@ class RGBPartNet(nn.Module): else: # evaluating f_c_, f_p_ = self.ae(x_c1_t2) x_c = self._decode_cano_feature(f_c_, n, t, device) - x_p = self._decode_pose_feature(f_p_, n, t, c, h, w, device) + x_p_ = self._decode_pose_feature(f_p_, n, t, c, h, w, device) + x_p = x_p_.view(n, t, self.pn_in_channels, self.h // 4, self.w // 4) return (x_c, x_p), None, None def _decode_appr_feature(self, f_a_, n, t, device): -- cgit v1.2.3