Add --max-sentence option for batching based on # sentences (f442f896) · Commits · Simon Will / fairseq

fairseq/criterions/cross_entropy.py

+3 −4

Changes for fairseq/criterions/cross_entropy.py: 3 added lines, 4 removed lines.
Original line number	Diff line number	Diff line
		@@ -14,9 +14,8 @@ from .fairseq_criterion import FairseqCriterion

		class CrossEntropyCriterion(FairseqCriterion):

		def __init__(self, padding_idx):
		super().__init__()
		self.padding_idx = padding_idx
		def __init__(self, args, dst_dict):
		super().__init__(args, dst_dict)

		def forward(self, model, sample):
		"""Compute the loss for the given sample.
		@@ -30,7 +29,7 @@ class CrossEntropyCriterion(FairseqCriterion):
		input = net_output.view(-1, net_output.size(-1))
		target = sample['target'].view(-1)
		loss = F.cross_entropy(input, target, size_average=False, ignore_index=self.padding_idx)
		sample_size = sample['ntokens']
		sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']
		logging_output = {
		'loss': loss.data[0],
		'sample_size': sample_size,

fairseq/criterions/fairseq_criterion.py

+3 −1

Changes for fairseq/criterions/fairseq_criterion.py: 3 added lines, 1 removed line.
Original line number	Diff line number	Diff line
		@@ -11,8 +11,10 @@ from torch.nn.modules.loss import _Loss

		class FairseqCriterion(_Loss):

		def __init__(self):
		def __init__(self, args, dst_dict):
		super().__init__()
		self.args = args
		self.padding_idx = dst_dict.pad()

		def forward(self, model, sample):
		"""Compute the loss for the given sample.

fairseq/criterions/label_smoothed_cross_entropy.py

+4 −5

Changes for fairseq/criterions/label_smoothed_cross_entropy.py: 4 added lines, 5 removed lines.
Original line number	Diff line number	Diff line
		@@ -43,10 +43,9 @@ class LabelSmoothedCrossEntropy(torch.autograd.Function):

		class LabelSmoothedCrossEntropyCriterion(FairseqCriterion):

		def __init__(self, eps, padding_idx=None, weights=None):
		super().__init__()
		self.eps = eps
		self.padding_idx = padding_idx
		def __init__(self, args, dst_dict, weights=None):
		super().__init__(args, dst_dict)
		self.eps = args.label_smoothing
		self.weights = weights

		def forward(self, model, sample):
		@@ -61,7 +60,7 @@ class LabelSmoothedCrossEntropyCriterion(FairseqCriterion):
		input = F.log_softmax(net_output.view(-1, net_output.size(-1)))
		target = sample['target'].view(-1)
		loss = LabelSmoothedCrossEntropy.apply(input, target, self.eps, self.padding_idx, self.weights)
		sample_size = sample['ntokens']
		sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']
		logging_output = {
		'loss': loss.data[0],
		'sample_size': sample_size,

fairseq/data.py

+45 −52

Changes for fairseq/data.py: 45 added lines, 52 removed lines.
Original line number	Diff line number	Diff line
		@@ -97,27 +97,26 @@ class LanguageDatasets(object):
		assert self.src_dict.unk() == self.dst_dict.unk()

		def train_dataloader(self, split, num_workers=0, max_tokens=None,
		max_positions=(1024, 1024), seed=None, epoch=1,
		sample_without_replacement=0,
		max_sentences=None, max_positions=(1024, 1024),
		seed=None, epoch=1, sample_without_replacement=0,
		sort_by_source_size=False):
		dataset = self.splits[split]
		with numpy_seed(seed):
		batch_sampler = shuffled_batches_by_size(
		dataset.src, dataset.dst,
		max_tokens=max_tokens, epoch=epoch,
		sample=sample_without_replacement,
		max_positions=max_positions,
		dataset.src, dataset.dst, max_tokens=max_tokens,
		max_sentences=max_sentences, epoch=epoch,
		sample=sample_without_replacement, max_positions=max_positions,
		sort_by_source_size=sort_by_source_size)
		return torch.utils.data.DataLoader(
		dataset, num_workers=num_workers, collate_fn=dataset.collater,
		batch_sampler=batch_sampler)

		def eval_dataloader(self, split, num_workers=0, batch_size=1,
		max_tokens=None, max_positions=(1024, 1024),
		def eval_dataloader(self, split, num_workers=0, max_tokens=None,
		max_sentences=None, max_positions=(1024, 1024),
		skip_invalid_size_inputs_valid_test=False):
		dataset = self.splits[split]
		batch_sampler = list(batches_by_size(
		dataset.src, dataset.dst, batch_size, max_tokens,
		dataset.src, dataset.dst, max_tokens, max_sentences,
		max_positions=max_positions,
		ignore_invalid_inputs=skip_invalid_size_inputs_valid_test))
		return torch.utils.data.DataLoader(
		@@ -220,29 +219,23 @@ def _valid_size(src_size, dst_size, max_positions):
		return True


		def batches_by_size(src, dst, batch_size=None, max_tokens=None,
		max_positions=(1024, 1024), ignore_invalid_inputs=False):
		"""Returns batches of indices sorted by size. Sequences with different
		source lengths are not allowed in the same batch."""
		assert isinstance(src, IndexedDataset) and isinstance(dst, IndexedDataset)
		if max_tokens is None:
		max_tokens = float('Inf')
		indices = np.argsort(src.sizes, kind='mergesort')

		def _make_batches(src, dst, indices, max_tokens, max_sentences, max_positions,
		ignore_invalid_inputs=False, allow_different_src_lens=False):
		batch = []

		def yield_batch(next_idx, num_tokens):
		if len(batch) == 0:
		return False
		if len(batch) == batch_size:
		if len(batch) == max_sentences:
		return True
		if src.sizes[batch[0]] != src.sizes[next_idx]:
		if num_tokens > max_tokens:
		return True
		if num_tokens >= max_tokens:
		if not allow_different_src_lens and \
		(src.sizes[batch[0]] != src.sizes[next_idx]):
		return True
		return False

		cur_max_size = 0
		sample_len = 0
		ignored = []
		for idx in indices:
		if not _valid_size(src.sizes[idx], dst.sizes[idx], max_positions):
		@@ -253,28 +246,48 @@ def batches_by_size(src, dst, batch_size=None, max_tokens=None,
		"Unable to handle input id {} of size {} / {}.".format(
		idx, src.sizes[idx], dst.sizes[idx]))

		if yield_batch(idx, cur_max_size * (len(batch) + 1)):
		sample_len = max(sample_len, src.sizes[idx], dst.sizes[idx])
		num_tokens = (len(batch) + 1) * sample_len
		if yield_batch(idx, num_tokens):
		yield batch
		batch = []
		cur_max_size = 0
		sample_len = max(src.sizes[idx], dst.sizes[idx])

		batch.append(idx)
		cur_max_size = max(cur_max_size, src.sizes[idx], dst.sizes[idx])

		if len(batch) > 0:
		yield batch

		if len(ignored) > 0:
		print("Warning! {} samples are either too short or too long "
		"and will be ignored, first few sample ids={}".format(len(ignored), ignored[:10]))

		if len(batch) > 0:
		yield batch

		def batches_by_size(src, dst, max_tokens=None, max_sentences=None,
		max_positions=(1024, 1024), ignore_invalid_inputs=False):
		"""Returns batches of indices sorted by size. Sequences with different
		source lengths are not allowed in the same batch."""
		assert isinstance(src, IndexedDataset) and isinstance(dst, IndexedDataset)
		if max_tokens is None:
		max_tokens = float('Inf')
		if max_sentences is None:
		max_sentences = float('Inf')
		indices = np.argsort(src.sizes, kind='mergesort')
		return _make_batches(
		src, dst, indices, max_tokens, max_sentences, max_positions,
		ignore_invalid_inputs, allow_different_src_lens=False)


		def shuffled_batches_by_size(src, dst, max_tokens=None, epoch=1, sample=0,
		max_positions=(1024, 1024), sort_by_source_size=False):
		def shuffled_batches_by_size(src, dst, max_tokens=None, max_sentences=None,
		epoch=1, sample=0, max_positions=(1024, 1024),
		sort_by_source_size=False):
		"""Returns batches of indices, bucketed by size and then shuffled. Batches
		may contain sequences of different lengths."""
		assert isinstance(src, IndexedDataset) and isinstance(dst, IndexedDataset)
		if max_tokens is None:
		max_tokens = float('Inf')
		if max_sentences is None:
		max_sentences = float('Inf')

		indices = np.random.permutation(len(src))

		@@ -282,30 +295,10 @@ def shuffled_batches_by_size(src, dst, max_tokens=None, epoch=1, sample=0,
		indices = indices[np.argsort(dst.sizes[indices], kind='mergesort')]
		indices = indices[np.argsort(src.sizes[indices], kind='mergesort')]

		def make_batches():
		batch = []
		sample_len = 0
		ignored = []
		for idx in indices:
		if not _valid_size(src.sizes[idx], dst.sizes[idx], max_positions):
		ignored.append(idx)
		continue
		sample_len = max(sample_len, src.sizes[idx], dst.sizes[idx])
		if len(batch) > 0 and (len(batch) + 1) * sample_len > max_tokens:
		yield batch
		batch = []
		sample_len = max(src.sizes[idx], dst.sizes[idx])

		batch.append(idx)

		if len(batch) > 0:
		yield batch

		if len(ignored) > 0:
		print("Warning! {} samples are either too short or too long "
		"and will be ignored, first few sample ids={}".format(len(ignored), ignored[:10]))
		batches = list(_make_batches(
		src, dst, indices, max_tokens, max_sentences, max_positions,
		ignore_invalid_inputs=True, allow_different_src_lens=True))

		batches = list(make_batches())
		if not sort_by_source_size:
		np.random.shuffle(batches)

fairseq/options.py

+3 −0

Changes for fairseq/options.py: 3 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -71,6 +71,9 @@ def add_optimization_args(parser):
		' dataset')
		group.add_argument('--curriculum', default=0, type=int, metavar='N',
		help='sort batches by source length for first N epochs')
		group.add_argument('--sentence-avg', action='store_true',
		help='normalize gradients by the number of sentences in a batch'
		' (default is to normalize by number of tokens)')
		return group