Reimplement cc_attention using einops.

[Leojc]

Motivation

Reimplement a faster and more readable cc_attention using einops.

Modification

1. Rewrite cc_attention.py
2. Add einops to the requirement.txt

[CLAassistant]

All committers have signed the CLA.

[ZwwWayne]

Hi @Leojc ，

1. Do you have some statistic numbers that verifies the current implementation?
2. I understand that the CUDA ops might be replaced, but the interfaces of this module should not be changed, e.g., in_channels should not be changed to in_dims.

[Leojc]

Hi @ZwwWayne ,
1.
(1) I calculate the time for one forward (input size is [2,512,97,97]). The CUDA version is 0.0299619 and my implementation is 0.00544943. So about 5x faster.
(2) I also test it in https://github.com/open-mmlab/mmsegmentation. I replace the cc_attention of ccnet in mmsegmentation. Then I run with a config of ccnet_r50-d8_769x769_40k_cityscapes.py. The CUDA version gets 78.46 mIoU, and this implementation gets 78.99.
2. Sorry about that! I just fixed it.

[ZwwWayne]

Thank you @Leojc , the statistical results are important. You may also highlight the speed comparison in the docstring of the module. I will invite @xvjiarui to have look.

[hellock]

Can it be implemented with torch.einsum so that we can omit a third-party dependency?

[Leojc]

Can it be implemented with torch.einsum so that we can omit a third-party dependency?

Yes, but that will decrease the readability greatly. Is that OK?

[Leojc]

Can it be implemented with torch.einsum so that we can omit a third-party dependency?

Yes, but that will decrease the readability greatly. Is that OK?

Wait... Maybe that won't decrease too much. I'll try implementing it.

[Leojc]

Think I should open a new PR.
Reimplement cc_attention using pure PyTorch #1201

[hellock]

Think I should open a new PR.
Reimplement cc_attention using pure PyTorch #1201

Looks great!