提交 8595d189 编写于 作者: W weishengyu

update format

上级 be806121
...@@ -23,7 +23,7 @@ from .. import dataloader ...@@ -23,7 +23,7 @@ from .. import dataloader
class MixDataset(Dataset): class MixDataset(Dataset):
def __init__(self, datasets_config): def __init__(self, datasets_config):
super(MixDataset, self).__init__() super().__init__()
self.dataset_list = [] self.dataset_list = []
start_idx = 0 start_idx = 0
end_idx = 0 end_idx = 0
......
...@@ -24,8 +24,9 @@ from ppcls.data import dataloader ...@@ -24,8 +24,9 @@ from ppcls.data import dataloader
class MixSampler(DistributedBatchSampler): class MixSampler(DistributedBatchSampler):
def __init__(self, dataset, batch_size, sample_configs, iter_per_epoch): def __init__(self, dataset, batch_size, sample_configs, iter_per_epoch):
super(MixSampler, self).__init__(dataset, batch_size) super().__init__(dataset, batch_size)
assert isinstance(dataset, MixDataset), "MixSampler only support MixDataset" assert isinstance(dataset,
MixDataset), "MixSampler only support MixDataset"
self.sampler_list = [] self.sampler_list = []
self.batch_size = batch_size self.batch_size = batch_size
self.start_list = [] self.start_list = []
...@@ -45,9 +46,11 @@ class MixSampler(DistributedBatchSampler): ...@@ -45,9 +46,11 @@ class MixSampler(DistributedBatchSampler):
assert batch_size_i <= len(dataset_list[i][2]) assert batch_size_i <= len(dataset_list[i][2])
config_i["batch_size"] = batch_size_i config_i["batch_size"] = batch_size_i
if sample_method == "DistributedBatchSampler": if sample_method == "DistributedBatchSampler":
sampler_i = DistributedBatchSampler(dataset_list[i][2], **config_i) sampler_i = DistributedBatchSampler(dataset_list[i][2],
**config_i)
else: else:
sampler_i = getattr(dataloader, sample_method)(dataset_list[i][2], **config_i) sampler_i = getattr(dataloader, sample_method)(
dataset_list[i][2], **config_i)
self.sampler_list.append(sampler_i) self.sampler_list.append(sampler_i)
self.iter_list.append(iter(sampler_i)) self.iter_list.append(iter(sampler_i))
self.length += len(dataset_list[i][2]) * ratio_i self.length += len(dataset_list[i][2]) * ratio_i
...@@ -62,7 +65,8 @@ class MixSampler(DistributedBatchSampler): ...@@ -62,7 +65,8 @@ class MixSampler(DistributedBatchSampler):
iter_i = iter(self.sampler_list[i]) iter_i = iter(self.sampler_list[i])
self.iter_list[i] = iter_i self.iter_list[i] = iter_i
batch_i = next(iter_i, None) batch_i = next(iter_i, None)
assert batch_i is not None, "dataset {} return None".format(i) assert batch_i is not None, "dataset {} return None".format(
i)
batch += [idx + self.start_list[i] for idx in batch_i] batch += [idx + self.start_list[i] for idx in batch_i]
if len(batch) == self.batch_size: if len(batch) == self.batch_size:
self.iter_counter += 1 self.iter_counter += 1
......
...@@ -42,7 +42,7 @@ class PKSampler(DistributedBatchSampler): ...@@ -42,7 +42,7 @@ class PKSampler(DistributedBatchSampler):
shuffle=True, shuffle=True,
drop_last=True, drop_last=True,
sample_method="sample_avg_prob"): sample_method="sample_avg_prob"):
super(PKSampler, self).__init__( super().__init__(
dataset, batch_size, shuffle=shuffle, drop_last=drop_last) dataset, batch_size, shuffle=shuffle, drop_last=drop_last)
assert batch_size % sample_per_id == 0, \ assert batch_size % sample_per_id == 0, \
"PKSampler configs error, Sample_per_id must be a divisor of batch_size." "PKSampler configs error, Sample_per_id must be a divisor of batch_size."
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册