Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import random | |
| import copy | |
| import time | |
| import warnings | |
| from torch.utils.data import Sampler | |
| from torch._six import int_classes as _int_classes | |
| # from configs.dog_breeds.dog_breed_class import get_partial_summary | |
| class TwoDatasetSampler(Sampler): | |
| """Wraps another sampler to yield a mini-batch of indices. | |
| Args: | |
| sampler (Sampler or Iterable): Base sampler. Can be any iterable object | |
| batch_size (int): Size of mini-batch. | |
| drop_last (bool): If ``True``, the sampler will drop the last batch if | |
| its size would be less than ``batch_size`` | |
| Example: | |
| >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False)) | |
| [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] | |
| >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True)) | |
| [[0, 1, 2], [3, 4, 5], [6, 7, 8]] | |
| """ | |
| def __init__(self, batch_size_half, size0, size1, shuffle=True, drop_last=True): | |
| # Since collections.abc.Iterable does not check for `__getitem__`, which | |
| # is one way for an object to be an iterable, we don't do an `isinstance` | |
| # check here. | |
| if not isinstance(batch_size_half, _int_classes) or isinstance(batch_size_half, bool) or \ | |
| batch_size_half <= 0: | |
| raise ValueError("batch_size should be a positive integer value, " | |
| "but got batch_size={}".format(batch_size_half*2)) | |
| if not isinstance(drop_last, bool): | |
| raise ValueError("drop_last should be a boolean value, but got " | |
| "drop_last={}".format(drop_last)) | |
| assert size0 >= size1 | |
| self.batch_size_half = batch_size_half | |
| self.size0 = size0 | |
| self.size1 = size1 | |
| self.shuffle = shuffle | |
| self.n_batches = self.size1//batch_size_half | |
| self.drop_last = drop_last | |
| def get_description(self): | |
| description = "\ | |
| This sampler samples equally from two different datasets" | |
| return description | |
| def __iter__(self): | |
| dataset0 = np.arange(self.size0) | |
| dataset1_init = np.arange(self.size1) + self.size0 | |
| if self.shuffle: | |
| np.random.shuffle(dataset0) | |
| dataset1 = [] | |
| for ind in range(self.size0 // self.size1 + 1): | |
| dataset1_part = dataset1_init.copy() | |
| if self.shuffle: | |
| np.random.shuffle(dataset1_part) | |
| dataset1.extend(dataset1_part) | |
| dataset0 = dataset0[0:self.n_batches*self.batch_size_half] | |
| dataset1 = dataset1[0:self.n_batches*self.batch_size_half] | |
| # import pdb; pdb.set_trace() | |
| for ind_batch in range(self.n_batches): | |
| d0 = dataset0[ind_batch*self.batch_size_half:(ind_batch+1)*self.batch_size_half] | |
| d1 = dataset1[ind_batch*self.batch_size_half:(ind_batch+1)*self.batch_size_half] | |
| batch = list(d0) + list(d1) | |
| # print(len(batch)) | |
| yield batch | |
| def __len__(self): | |
| # Can only be called if self.sampler has __len__ implemented | |
| # We cannot enforce this condition, so we turn off typechecking for the | |
| # implementation below. | |
| # Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] | |
| '''if self.drop_last: | |
| return len(self.sampler) // self.batch_size # type: ignore | |
| else: | |
| return (len(self.sampler) + self.batch_size - 1) // self.batch_size # type: ignore''' | |
| return self.n_batches | |