[Fix] fix the wrong weight reference bug in BaseTransformerLayer

[gaotongxiao]

Fixes the wrong function/weight reference bug in BaseTransformerLayer when batch_first is True.

Motivation

There are some cases that users clone a model with copy.deepcopy to create, for example, a ModuleList. However, things become tricky when it comes to BaseTransformerLayer. If we create a ModuleList on GPU in this way with batch_first=True and run the forward function, an error will prompt saying the weights are not on GPU.

baselayer = BaseTransformerLayer(
    operation_order=('self_attn', 'ffn'),
    batch_first=True,
    attn_cfgs=dict(
        type='MultiheadAttention',
        embed_dims=256,
        num_heads=8,
    ),
)
baselayers = ModuleList([copy.deepcopy(baselayer) for _ in range(2)])
baselayers.to('cuda')
x = torch.rand(2, 10, 256).cuda()
out = baselayers[0](x)

The error:

RuntimeError: Tensor for argument #3 'mat2' is on CPU, but expected it to be on GPU (while checking arguments for addmm)

It turned out to be the problem of forward_wrapper, which somehow stores the function pointer of self.attn.forward during the initialization and would always call this address after transpose operations. Such logic works well for this module but not for modules deepcopied from it, which will still call the forward function of the original module. Essentially they are still using the weights of the original module to conduct computations!

mmcv/mmcv/cnn/bricks/transformer.py

Lines 105 to 125 in e8489a7

	if self.batch_first:
	def _bnc_to_nbc(forward):
	"""Because the dataflow('key', 'query', 'value') of
	``torch.nn.MultiheadAttention`` is (num_query, batch,
	embed_dims), We should adjust the shape of dataflow from
	batch_first (batch, num_query, embed_dims) to num_query_first
	(num_query ,batch, embed_dims), and recover ``attn_output``
	from num_query_first to batch_first."""
	def forward_wrapper(**kwargs):
	convert_keys = ('key', 'query', 'value')
	for key in kwargs.keys():
	if key in convert_keys:
	kwargs[key] = kwargs[key].transpose(0, 1)
	attn_output, attn_output_weights = forward(**kwargs)
	return attn_output.transpose(0, 1), attn_output_weights
	return forward_wrapper
	self.attn.forward = _bnc_to_nbc(self.attn.forward)

Modification

Hardcoding a function pointer in the initializer is not a good practice. I move the transpose logic to forward and it should be much safer. A unit test for this case is also added.

BC-breaking (Optional)

No

[jshilong]

LGTM