Pytorch Implementation of CoAtNet, the SOTA image classification model described in https://arxiv.org/abs/2106.04803.
CoatNet.py contains the coatnet class and accompanying initialization functions for each size of coatnet.
The translation equivariant weight is produced within the code by taking advantage of matrix multiplication's inherent order preserving nature.
With the equation pictured above for finding the index of the respective equivariant weight for an index of the K*Q number being i-j, the entirety of the indexing for this problem can be solved with a diagonally similar matrix of equivarant weight index. For example, if the equivarant weight was (1, 2, 3, 4, 5, 6, 7), the respective additive for K*Q would be:
This result may seem odd, but really is not when one looks at the respective index for itself: