提交 cb9c6fad 编写于 作者: X xiefangqi

fix numpyslice issue

上级 7cb567eb
......@@ -3219,33 +3219,9 @@ class GeneratorDataset(MappableDataset):
def __init__(self, source, column_names=None, column_types=None, schema=None, num_samples=None,
num_parallel_workers=1, shuffle=None, sampler=None, num_shards=None, shard_id=None):
super().__init__(num_parallel_workers)
self.source = source
self.sampler = _select_sampler(num_samples, sampler, shuffle, num_shards, shard_id)
if self.sampler is not None and hasattr(source, "__getitem__"):
if isinstance(self.sampler, (samplers.SequentialSampler, samplers.DistributedSampler,
samplers.RandomSampler, samplers.SubsetRandomSampler,
samplers.WeightedRandomSampler, samplers.Sampler)):
sampler_instance = self.sampler.create()
sampler_instance.set_num_rows(len(source))
sampler_instance.initialize()
if num_parallel_workers > 1:
self.source = (lambda: _cpp_sampler_fn_mp(sampler_instance, source, num_parallel_workers))
else:
self.source = (lambda: _cpp_sampler_fn(sampler_instance, source))
else:
if num_parallel_workers > 1:
self.source = (lambda: _py_sampler_fn_mp(self.sampler, num_samples, source, num_parallel_workers))
else:
self.source = (lambda: _py_sampler_fn(self.sampler, num_samples, source))
else:
try:
iter(source)
except TypeError:
# Use generator function if input callable
self.source = (lambda: _generator_fn(source, num_samples))
else:
# Use iterator function if input is iterable
# Random accessible input is also iterable
self.source = (lambda: _iter_fn(source, num_samples))
self.num_samples = num_samples
if column_names is not None and not isinstance(column_names, list):
column_names = [column_names]
......@@ -3310,9 +3286,35 @@ class GeneratorDataset(MappableDataset):
new_op.num_parallel_workers = copy.deepcopy(self.num_parallel_workers, memodict)
new_op.column_types = copy.deepcopy(self.column_types, memodict)
new_op.column_names = copy.deepcopy(self.column_names, memodict)
new_op.source = self.source
new_op.sampler = self.sampler
new_op.num_samples = copy.deepcopy(self.num_samples, memodict)
new_op.sampler = copy.deepcopy(self.sampler)
if new_op.sampler is not None and hasattr(self.source, "__getitem__"):
if isinstance(new_op.sampler, (samplers.SequentialSampler, samplers.DistributedSampler,
samplers.RandomSampler, samplers.SubsetRandomSampler,
samplers.WeightedRandomSampler, samplers.Sampler)):
sampler_instance = new_op.sampler.create()
sampler_instance.set_num_rows(len(self.source))
sampler_instance.initialize()
if new_op.num_parallel_workers > 1:
new_op.source = (lambda: _cpp_sampler_fn_mp(sampler_instance, self.source, new_op.num_parallel_workers))
else:
new_op.source = (lambda: _cpp_sampler_fn(sampler_instance, self.source))
else:
if new_op.num_parallel_workers > 1:
new_op.source = (lambda: _py_sampler_fn_mp(new_op.sampler, new_op.num_samples, self.source, new_op.num_parallel_workers))
else:
new_op.source = (lambda: _py_sampler_fn(new_op.sampler, new_op.num_samples, self.source))
else:
try:
iter(self.source)
except TypeError:
# Use generator function if input callable
new_op.source = (lambda: _generator_fn(self.source, new_op.num_samples))
else:
# Use iterator function if input is iterable
# Random accessible input is also iterable
new_op.source = (lambda: _iter_fn(self.source, new_op.num_samples))
return new_op
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册