Skip to content

Commit

Permalink
make MultipleDatasetIterator only if necessary (OpenNMT#1436)
Browse files Browse the repository at this point in the history
  • Loading branch information
francoishernandez authored and vince62s committed May 17, 2019
1 parent 87e967d commit fc6ec53
Showing 1 changed file with 7 additions and 8 deletions.
15 changes: 7 additions & 8 deletions onmt/train_single.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,15 +99,14 @@ def main(opt, device_id):
opt, device_id, model, fields, optim, model_saver=model_saver)

train_iterables = []
for train_id in opt.data_ids:
if train_id:
if len(opt.data_ids) > 1:
for train_id in opt.data_ids:
shard_base = "train_" + train_id
else:
shard_base = "train"
iterable = build_dataset_iter(shard_base, fields, opt, multi=True)
train_iterables.append(iterable)

train_iter = MultipleDatasetIterator(train_iterables, device_id, opt)
iterable = build_dataset_iter(shard_base, fields, opt, multi=True)
train_iterables.append(iterable)
train_iter = MultipleDatasetIterator(train_iterables, device_id, opt)
else:
train_iter = build_dataset_iter("train", fields, opt)

valid_iter = build_dataset_iter(
"valid", fields, opt, is_train=False)
Expand Down

0 comments on commit fc6ec53

Please sign in to comment.