Inefficient _layer_norm implementation in layerwise_attention.py

[K024]

🐛 Bug

Iterations along batch_size axis are slow.

Code:

COMET/comet/modules/layerwise_attention.py

Lines 96 to 122 in 26d5779

	def _layer_norm(tensor, broadcast_mask, mask):
	tensor_masked = tensor * broadcast_mask
	batch_size, _, input_dim = tensors[0].size()
	num_elements_not_masked = torch.tensor(
	[mask[i].sum() * input_dim for i in range(batch_size)],
	device=tensor.device,
	)
	# mean for each sentence
	mean = torch.sum(torch.sum(tensor_masked, dim=2), dim=1)
	mean = mean / num_elements_not_masked
	variance = torch.vstack(
	[
	torch.sum(((tensor_masked[i] - mean[i]) * broadcast_mask[i]) ** 2)
	/ num_elements_not_masked[i]
	for i in range(batch_size)
	]
	)
	normalized_tensor = torch.vstack(
	[
	((tensor[i] - mean[i]) / torch.sqrt(variance[i] + 1e-12)).unsqueeze(
	0
	)
	for i in range(batch_size)
	]
	)
	return normalized_tensor

This also prevents wmt20-comet-qe-da to be exported to onnx format.

To Reproduce

N/A

Expected behaviour

Use .view to create a proper shape for broadcasting.

Preferred code:

        def _layer_norm(tensor, broadcast_mask, mask):
            tensor_masked = tensor * broadcast_mask
            batch_size, _, input_dim = tensors[0].size()

            # mean for each sentence
            num_elements_not_masked = mask.sum(1) * input_dim
            mean = tensor_masked.view(batch_size, -1).sum(1)
            mean = (mean / num_elements_not_masked).view(batch_size, 1, 1)

            variance = (((tensor_masked - mean) * broadcast_mask) ** 2).view(batch_size, -1).sum(1) / num_elements_not_masked
            normalized_tensor = (tensor - mean) / torch.sqrt(variance + 1e-12).view(batch_size, 1, 1)
            return normalized_tensor

Screenshots

N/A

Environment

N/A

Additional context

N/A

[ricardorei]

This is a good suggestion. I tested it and everything worked well its currently in master.