From 50eedae4f320c446544772fb2b0abcbce1be7590 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/auto_encoder.py | 2 +- models/model.py | 20 ++++++++++++++--- models/rgb_part_net.py | 20 +++-------------- utils/triplet_loss.py | 58 +++++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 71 insertions(+), 29 deletions(-) diff --git a/models/auto_encoder.py b/models/auto_encoder.py index 2d715db..1ef7494 100644 --- a/models/auto_encoder.py +++ b/models/auto_encoder.py @@ -171,7 +171,7 @@ class AutoEncoder(nn.Module): return ( (f_a_c1_t2_, f_c_c1_t2_, f_p_c1_t2_), - (xrecon_loss, cano_cons_loss, pose_sim_loss * 10) + torch.stack((xrecon_loss, cano_cons_loss, pose_sim_loss * 10)) ) else: # evaluating return f_c_c1_t2_, f_p_c1_t2_ diff --git a/models/model.py b/models/model.py index 82d6461..5aa6436 100644 --- a/models/model.py +++ b/models/model.py @@ -18,6 +18,7 @@ from utils.configuration import DataloaderConfiguration, \ SystemConfiguration from utils.dataset import CASIAB, ClipConditions, ClipViews, ClipClasses from utils.sampler import TripletSampler +from utils.triplet_loss import JointBatchAllTripletLoss class Model: @@ -67,6 +68,7 @@ class Model: self._dataset_sig: str = 'undefined' self.rgb_pn: Optional[RGBPartNet] = None + self.ba_triplet_loss: Optional[JointBatchAllTripletLoss] = None self.optimizer: Optional[optim.Adam] = None self.scheduler: Optional[optim.lr_scheduler.StepLR] = None self.writer: Optional[SummaryWriter] = None @@ -140,7 +142,8 @@ class Model: dataset = self._parse_dataset_config(dataset_config) dataloader = self._parse_dataloader_config(dataset, dataloader_config) # Prepare for model, optimizer and scheduler - model_hp = self.hp.get('model', {}) + model_hp: dict = self.hp.get('model', {}).copy() + triplet_margins = model_hp.pop('triplet_margins', (0.2, 0.2)) optim_hp: dict = self.hp.get('optimizer', {}).copy() start_iter = optim_hp.pop('start_iter', 0) ae_optim_hp = optim_hp.pop('auto_encoder', {}) @@ -150,6 +153,9 @@ class Model: sched_hp = self.hp.get('scheduler', {}) self.rgb_pn = RGBPartNet(self.in_channels, self.in_size, **model_hp, image_log_on=self.image_log_on) + self.ba_triplet_loss = JointBatchAllTripletLoss( + self.rgb_pn.hpm_num_parts, triplet_margins + ) # Try to accelerate computation using CUDA or others self.rgb_pn = self.rgb_pn.to(self.device) self.optimizer = optim.Adam([ @@ -193,10 +199,18 @@ class Model: # forward + backward + optimize x_c1 = batch_c1['clip'].to(self.device) x_c2 = batch_c2['clip'].to(self.device) + feature, ae_losses, images = self.rgb_pn(x_c1, x_c2) 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) + y = y.repeat(self.rgb_pn.num_total_parts, 1) + triplet_loss = self.ba_triplet_loss(feature, y) + losses = torch.cat(( + ae_losses, + torch.stack(( + triplet_loss[:self.rgb_pn.hpm_num_parts].mean(), + triplet_loss[self.rgb_pn.hpm_num_parts:].mean() + )) + )) loss = losses.sum() loss.backward() self.optimizer.step() 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) diff --git a/utils/triplet_loss.py b/utils/triplet_loss.py index 954def2..0df2188 100644 --- a/utils/triplet_loss.py +++ b/utils/triplet_loss.py @@ -11,6 +11,25 @@ class BatchAllTripletLoss(nn.Module): def forward(self, x, y): p, n, c = x.size() + dist = self._batch_distance(x) + positive_negative_dist = self._hard_distance(dist, y, p, n) + all_loss = F.relu(self.margin + positive_negative_dist).view(p, -1) + parted_loss_mean = self._none_zero_parted_mean(all_loss) + + return parted_loss_mean + + @staticmethod + def _hard_distance(dist, y, p, n): + hard_positive_mask = y.unsqueeze(1) == y.unsqueeze(2) + hard_negative_mask = y.unsqueeze(1) != y.unsqueeze(2) + all_hard_positive = dist[hard_positive_mask].view(p, n, -1, 1) + all_hard_negative = dist[hard_negative_mask].view(p, n, 1, -1) + positive_negative_dist = all_hard_positive - all_hard_negative + + return positive_negative_dist + + @staticmethod + def _batch_distance(x): # Euclidean distance p x n x n x_squared_sum = torch.sum(x ** 2, dim=2) x1_squared_sum = x_squared_sum.unsqueeze(2) @@ -20,17 +39,40 @@ class BatchAllTripletLoss(nn.Module): F.relu(x1_squared_sum - 2 * x1_times_x2_sum + x2_squared_sum) ) - hard_positive_mask = y.unsqueeze(1) == y.unsqueeze(2) - hard_negative_mask = y.unsqueeze(1) != y.unsqueeze(2) - all_hard_positive = dist[hard_positive_mask].view(p, n, -1, 1) - all_hard_negative = dist[hard_negative_mask].view(p, n, 1, -1) - positive_negative_dist = all_hard_positive - all_hard_negative - all_loss = F.relu(self.margin + positive_negative_dist).view(p, -1) + return dist + @staticmethod + def _none_zero_parted_mean(all_loss): # Non-zero parted mean 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.mean() - return loss + return parted_loss_mean + + +class JointBatchAllTripletLoss(BatchAllTripletLoss): + def __init__( + self, + hpm_num_parts: int, + margins: tuple[float, float] = (0.2, 0.2) + ): + super().__init__() + self.hpm_num_parts = hpm_num_parts + self.margin_hpm, self.margin_pn = margins + + def forward(self, x, y): + p, n, c = x.size() + + dist = self._batch_distance(x) + positive_negative_dist = self._hard_distance(dist, y, p, n) + hpm_part_loss = F.relu( + self.margin_hpm + positive_negative_dist[:self.hpm_num_parts] + ).view(self.hpm_num_parts, -1) + pn_part_loss = F.relu( + self.margin_pn + positive_negative_dist[self.hpm_num_parts:] + ).view(p - self.hpm_num_parts, -1) + all_loss = torch.cat((hpm_part_loss, pn_part_loss)).view(p, -1) + parted_loss_mean = self._none_zero_parted_mean(all_loss) + + return parted_loss_mean -- 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 --- config.py | 12 ++---------- models/auto_encoder.py | 4 ++-- models/part_net.py | 18 ++++-------------- models/rgb_part_net.py | 37 ++++++++++++++++++------------------- 4 files changed, 26 insertions(+), 45 deletions(-) diff --git a/config.py b/config.py index 424bf5b..88ad371 100644 --- a/config.py +++ b/config.py @@ -49,22 +49,14 @@ config: Configuration = { # Auto-encoder feature channels coefficient 'ae_feature_channels': 64, # Appearance, canonical and pose feature dimensions - 'f_a_c_p_dims': (128, 128, 64), + 'f_a_c_p_dims': (192, 192, 96), # Use 1x1 convolution in dimensionality reduction 'hpm_use_1x1conv': False, # HPM pyramid scales, of which sum is number of parts 'hpm_scales': (1, 2, 4), # Global pooling method 'hpm_use_avg_pool': True, - 'hpm_use_max_pool': False, - # FConv feature channels coefficient - 'fpfe_feature_channels': 32, - # FConv blocks kernel sizes - 'fpfe_kernel_sizes': ((5, 3), (3, 3), (3, 3)), - # FConv blocks paddings - 'fpfe_paddings': ((2, 1), (1, 1), (1, 1)), - # FConv blocks halving - 'fpfe_halving': (0, 2, 3), + 'hpm_use_max_pool': True, # Attention squeeze ratio 'tfa_squeeze_ratio': 4, # Number of parts after Part Net diff --git a/models/auto_encoder.py b/models/auto_encoder.py index 1ef7494..e6a3e60 100644 --- a/models/auto_encoder.py +++ b/models/auto_encoder.py @@ -106,14 +106,14 @@ class Decoder(nn.Module): self.trans_conv4 = DCGANConvTranspose2d(feature_channels, out_channels, is_last_layer=True) - def forward(self, f_appearance, f_canonical, f_pose, cano_only=False): + def forward(self, f_appearance, f_canonical, f_pose, is_feature_map=False): x = torch.cat((f_appearance, f_canonical, f_pose), dim=1) x = self.fc(x) x = x.view(-1, self.feature_channels * 8, self.h_0, self.w_0) x = F.relu(x, inplace=True) x = self.trans_conv1(x) x = self.trans_conv2(x) - if cano_only: + if is_feature_map: return x x = self.trans_conv3(x) x = torch.sigmoid(self.trans_conv4(x)) diff --git a/models/part_net.py b/models/part_net.py index 62a2bac..29cf9cd 100644 --- a/models/part_net.py +++ b/models/part_net.py @@ -110,32 +110,22 @@ class TemporalFeatureAggregator(nn.Module): class PartNet(nn.Module): def __init__( self, - in_channels: int = 3, - feature_channels: int = 32, - kernel_sizes: tuple[tuple, ...] = ((5, 3), (3, 3), (3, 3)), - paddings: tuple[tuple, ...] = ((2, 1), (1, 1), (1, 1)), - halving: tuple[int, ...] = (0, 2, 3), + in_channels: int = 128, squeeze_ratio: int = 4, num_part: int = 16 ): super().__init__() self.num_part = num_part - self.fpfe = FrameLevelPartFeatureExtractor( - in_channels, feature_channels, kernel_sizes, paddings, halving - ) - - num_fconv_blocks = len(self.fpfe.fconv_blocks) - self.tfa_in_channels = feature_channels * 2 ** (num_fconv_blocks - 1) self.tfa = TemporalFeatureAggregator( - self.tfa_in_channels, squeeze_ratio, self.num_part + in_channels, squeeze_ratio, self.num_part ) self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) def forward(self, x): - n, t, _, _, _ = x.size() - x = self.fpfe(x) + n, t, c, h, w = x.size() + x = x.view(n * t, c, h, w) # n * t x c x h x w # Horizontal Pooling 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/model.py | 3 ++- models/rgb_part_net.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/models/model.py b/models/model.py index 5899fc0..90d48e0 100644 --- a/models/model.py +++ b/models/model.py @@ -314,7 +314,8 @@ class Model: ) # Init models - model_hp = self.hp.get('model', {}) + model_hp: dict = self.hp.get('model', {}).copy() + model_hp.pop('triplet_margins', None) self.rgb_pn = RGBPartNet(self.in_channels, self.in_size, **model_hp) # Try to accelerate computation using CUDA or others self.rgb_pn = self.rgb_pn.to(self.device) 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 From 9001f7e13d8985b220bd218d8de716bc586dbdcf Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Fri, 26 Feb 2021 20:17:03 +0800 Subject: Update default config --- config.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/config.py b/config.py index 88ad371..03f2f0d 100644 --- a/config.py +++ b/config.py @@ -37,7 +37,7 @@ config: Configuration = { # Batch size (pr, k) # `pr` denotes number of persons # `k` denotes number of sequences per person - 'batch_size': (4, 8), + 'batch_size': (4, 6), # Number of workers of Dataloader 'num_workers': 4, # Faster data transfer from RAM to GPU if enabled @@ -64,7 +64,7 @@ config: Configuration = { # Embedding dimension for each part 'embedding_dims': 256, # Triplet loss margins for HPM and PartNet - 'triplet_margins': (0.2, 0.2), + 'triplet_margins': (1.5, 1.5), }, 'optimizer': { # Global parameters @@ -83,15 +83,15 @@ config: Configuration = { # 'amsgrad': False, # Local parameters (override global ones) - 'auto_encoder': { - 'weight_decay': 0.001 - }, + # 'auto_encoder': { + # 'weight_decay': 0.001 + # }, }, 'scheduler': { # Period of learning rate decay 'step_size': 500, # Multiplicative factor of decay - 'gamma': 0.9, + 'gamma': 1, } }, # Model metadata -- cgit v1.2.3