diff options
author | Jordan Gong <jordan.gong@protonmail.com> | 2021-03-14 21:08:14 +0800 |
---|---|---|
committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-03-14 21:08:14 +0800 |
commit | f953000c5b9490c040b68cbc233518c756ac00ab (patch) | |
tree | 085742fa2a52524bb29529d9d254ac80735ada66 | |
parent | c97bed6721fc1a7186719fef0c340d6d03057891 (diff) | |
parent | 17ae5e529475a7e47fdde0ce69982b5dfe34f2eb (diff) |
Merge branch 'python3.8' into python3.7
-rw-r--r-- | config.py | 6 | ||||
-rw-r--r-- | models/rgb_part_net.py | 2 | ||||
-rw-r--r-- | utils/sampler.py | 35 |
3 files changed, 35 insertions, 8 deletions
@@ -92,9 +92,9 @@ config = { '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 6e72e38..fba3485 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -59,7 +59,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() diff --git a/utils/sampler.py b/utils/sampler.py index 0977f94..581d7a2 100644 --- a/utils/sampler.py +++ b/utils/sampler.py @@ -15,7 +15,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 @@ -26,15 +37,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 |