From 263b8001ca1b25a43d1c87f187423054e141925d Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 14 Mar 2021 21:07:03 +0800 Subject: Fix unbalanced datasets --- utils/sampler.py | 35 +++++++++++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/utils/sampler.py b/utils/sampler.py index cdf1984..0c9872c 100644 --- a/utils/sampler.py +++ b/utils/sampler.py @@ -16,7 +16,18 @@ class TripletSampler(data.Sampler): ): super().__init__(data_source) self.metadata_labels = data_source.metadata['labels'] + metadata_conditions = data_source.metadata['conditions'] + self.subsets = {} + for condition in metadata_conditions: + pre, _ = condition.split('-') + if self.subsets.get(pre, None) is None: + self.subsets[pre] = [] + self.subsets[pre].append(condition) + self.num_subsets = len(self.subsets) + self.num_seq = {pre: len(seq) for (pre, seq) in self.subsets.items()} + self.min_num_seq = min(self.num_seq.values()) self.labels = data_source.labels + self.conditions = data_source.conditions self.length = len(self.labels) self.indexes = np.arange(0, self.length) (self.pr, self.k) = batch_size @@ -27,15 +38,31 @@ class TripletSampler(data.Sampler): # Sample pr subjects by sampling labels appeared in dataset sampled_subjects = random.sample(self.metadata_labels, k=self.pr) for label in sampled_subjects: - clips_from_subject = self.indexes[self.labels == label].tolist() + mask = self.labels == label + # Fix unbalanced datasets + if self.num_subsets > 1: + condition_mask = np.zeros(self.conditions.shape, dtype=bool) + for num, conditions_ in zip( + self.num_seq.values(), self.subsets.values() + ): + if num > self.min_num_seq: + conditions = random.sample( + conditions_, self.min_num_seq + ) + else: + conditions = conditions_ + for condition in conditions: + condition_mask |= self.conditions == condition + mask &= condition_mask + clips = self.indexes[mask].tolist() # Sample k clips from the subject without replacement if # have enough clips, k more clips will sampled for # disentanglement k = self.k * 2 - if len(clips_from_subject) >= k: - _sampled_indexes = random.sample(clips_from_subject, k=k) + if len(clips) >= k: + _sampled_indexes = random.sample(clips, k=k) else: - _sampled_indexes = random.choices(clips_from_subject, k=k) + _sampled_indexes = random.choices(clips, k=k) sampled_indexes += _sampled_indexes yield sampled_indexes -- cgit v1.2.3 From da922be042d96338a3f207386e410b6746d046f5 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 14 Mar 2021 21:07:28 +0800 Subject: Bug fix when transforming and new config --- config.py | 6 +++--- models/rgb_part_net.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/config.py b/config.py index 8abeba3..c928067 100644 --- a/config.py +++ b/config.py @@ -94,9 +94,9 @@ config: Configuration = { 'final_gamma': 0.01, # Local parameters (override global ones) - 'hpm': { - 'final_gamma': 0.001 - } + # 'hpm': { + # 'final_gamma': 0.001 + # } } }, # Model metadata diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index c38a567..c136040 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -57,7 +57,7 @@ class RGBPartNet(nn.Module): if self.training: return x_c, x_p, ae_losses, images else: - return torch.cat((x_c, x_p)).unsqueeze(1).view(-1) + return torch.cat((x_c.view(-1), x_p.view(-1))) def _disentangle(self, x_c1_t2, x_c2_t2=None): n, t, c, h, w = x_c1_t2.size() -- cgit v1.2.3