release code

Dassl.ProGrad.pytorch/dassl/data/samplers.py

@@ -0,0 +1,205 @@
+import copy
+import numpy as np
+import random
+from collections import defaultdict
+from torch.utils.data.sampler import Sampler, RandomSampler, SequentialSampler
+
+
+class RandomDomainSampler(Sampler):
+    """Randomly samples N domains each with K images
+    to form a minibatch of size N*K.
+
+    Args:
+        data_source (list): list of Datums.
+        batch_size (int): batch size.
+        n_domain (int): number of domains to sample in a minibatch.
+    """
+
+    def __init__(self, data_source, batch_size, n_domain):
+        self.data_source = data_source
+
+        # Keep track of image indices for each domain
+        self.domain_dict = defaultdict(list)
+        for i, item in enumerate(data_source):
+            self.domain_dict[item.domain].append(i)
+        self.domains = list(self.domain_dict.keys())
+
+        # Make sure each domain has equal number of images
+        if n_domain is None or n_domain <= 0:
+            n_domain = len(self.domains)
+        assert batch_size % n_domain == 0
+        self.n_img_per_domain = batch_size // n_domain
+
+        self.batch_size = batch_size
+        # n_domain denotes number of domains sampled in a minibatch
+        self.n_domain = n_domain
+        self.length = len(list(self.__iter__()))
+
+    def __iter__(self):
+        domain_dict = copy.deepcopy(self.domain_dict)
+        final_idxs = []
+        stop_sampling = False
+
+        while not stop_sampling:
+            selected_domains = random.sample(self.domains, self.n_domain)
+
+            for domain in selected_domains:
+                idxs = domain_dict[domain]
+                selected_idxs = random.sample(idxs, self.n_img_per_domain)
+                final_idxs.extend(selected_idxs)
+
+                for idx in selected_idxs:
+                    domain_dict[domain].remove(idx)
+
+                remaining = len(domain_dict[domain])
+                if remaining < self.n_img_per_domain:
+                    stop_sampling = True
+
+        return iter(final_idxs)
+
+    def __len__(self):
+        return self.length
+
+
+class SeqDomainSampler(Sampler):
+    """Sequential domain sampler, which randomly samples K
+    images from each domain to form a minibatch.
+
+    Args:
+        data_source (list): list of Datums.
+        batch_size (int): batch size.
+    """
+
+    def __init__(self, data_source, batch_size):
+        self.data_source = data_source
+
+        # Keep track of image indices for each domain
+        self.domain_dict = defaultdict(list)
+        for i, item in enumerate(data_source):
+            self.domain_dict[item.domain].append(i)
+        self.domains = list(self.domain_dict.keys())
+        self.domains.sort()
+
+        # Make sure each domain has equal number of images
+        n_domain = len(self.domains)
+        assert batch_size % n_domain == 0
+        self.n_img_per_domain = batch_size // n_domain
+
+        self.batch_size = batch_size
+        # n_domain denotes number of domains sampled in a minibatch
+        self.n_domain = n_domain
+        self.length = len(list(self.__iter__()))
+
+    def __iter__(self):
+        domain_dict = copy.deepcopy(self.domain_dict)
+        final_idxs = []
+        stop_sampling = False
+
+        while not stop_sampling:
+            for domain in self.domains:
+                idxs = domain_dict[domain]
+                selected_idxs = random.sample(idxs, self.n_img_per_domain)
+                final_idxs.extend(selected_idxs)
+
+                for idx in selected_idxs:
+                    domain_dict[domain].remove(idx)
+
+                remaining = len(domain_dict[domain])
+                if remaining < self.n_img_per_domain:
+                    stop_sampling = True
+
+        return iter(final_idxs)
+
+    def __len__(self):
+        return self.length
+
+
+class RandomClassSampler(Sampler):
+    """Randomly samples N classes each with K instances to
+    form a minibatch of size N*K.
+
+    Modified from https://github.com/KaiyangZhou/deep-person-reid.
+
+    Args:
+        data_source (list): list of Datums.
+        batch_size (int): batch size.
+        n_ins (int): number of instances per class to sample in a minibatch.
+    """
+
+    def __init__(self, data_source, batch_size, n_ins):
+        if batch_size < n_ins:
+            raise ValueError(
+                "batch_size={} must be no less "
+                "than n_ins={}".format(batch_size, n_ins)
+            )
+
+        self.data_source = data_source
+        self.batch_size = batch_size
+        self.n_ins = n_ins
+        self.ncls_per_batch = self.batch_size // self.n_ins
+        self.index_dic = defaultdict(list)
+        for index, item in enumerate(data_source):
+            self.index_dic[item.label].append(index)
+        self.labels = list(self.index_dic.keys())
+        assert len(self.labels) >= self.ncls_per_batch
+
+        # estimate number of images in an epoch
+        self.length = len(list(self.__iter__()))
+
+    def __iter__(self):
+        batch_idxs_dict = defaultdict(list)
+
+        for label in self.labels:
+            idxs = copy.deepcopy(self.index_dic[label])
+            if len(idxs) < self.n_ins:
+                idxs = np.random.choice(idxs, size=self.n_ins, replace=True)
+            random.shuffle(idxs)
+            batch_idxs = []
+            for idx in idxs:
+                batch_idxs.append(idx)
+                if len(batch_idxs) == self.n_ins:
+                    batch_idxs_dict[label].append(batch_idxs)
+                    batch_idxs = []
+
+        avai_labels = copy.deepcopy(self.labels)
+        final_idxs = []
+
+        while len(avai_labels) >= self.ncls_per_batch:
+            selected_labels = random.sample(avai_labels, self.ncls_per_batch)
+            for label in selected_labels:
+                batch_idxs = batch_idxs_dict[label].pop(0)
+                final_idxs.extend(batch_idxs)
+                if len(batch_idxs_dict[label]) == 0:
+                    avai_labels.remove(label)
+
+        return iter(final_idxs)
+
+    def __len__(self):
+        return self.length
+
+
+def build_sampler(
+    sampler_type,
+    cfg=None,
+    data_source=None,
+    batch_size=32,
+    n_domain=0,
+    n_ins=16
+):
+    if sampler_type == "RandomSampler":
+        return RandomSampler(data_source)
+
+    elif sampler_type == "SequentialSampler":
+        return SequentialSampler(data_source)
+
+    elif sampler_type == "RandomDomainSampler":
+        return RandomDomainSampler(data_source, batch_size, n_domain)
+
+    elif sampler_type == "SeqDomainSampler":
+        return SeqDomainSampler(data_source, batch_size)
+
+    elif sampler_type == "RandomClassSampler":
+        return RandomClassSampler(data_source, batch_size, n_ins)
+
+    else:
+        raise ValueError("Unknown sampler type: {}".format(sampler_type))