adapt roll op for swin transformer

ppcls/arch/backbone/legendary_models/swin_transformer.py

@@ -42,6 +42,50 @@ MODEL_URLS = {
 
 __all__ = list(MODEL_URLS.keys())
 
+class RollWithIndexSelect(paddle.autograd.PyLayer):
+    @staticmethod
+    def forward(ctx, input1, index_fp, index_bp):
+        N, H, W, C = input1.shape
+        ctx.input1 = input1
+        ctx.index_bp = index_bp
+        result = input1.reshape([N, H * W, C]).index_select(index_fp, 1).reshape([N, H, W, C])
+        return result
+
+    @staticmethod
+    def backward(ctx, grad):
+        input1 = ctx.input1
+        N, H, W, C = input1.shape
+        index_bp = ctx.index_bp
+        grad_input = grad.reshape([N, H * W, C]).index_select(index_bp, 1).reshape([N, H, W, C])
+        return grad_input, None, None
+
+roll_with_index_select = RollWithIndexSelect.apply
+
+def get_roll_index(H, W, shifts, place):
+    index = paddle.arange(0, H * W, dtype='int64').reshape([H, W]) # cpu
+    index_fp = paddle.roll(index, shifts=shifts, axis=(0, 1)).reshape([-1]) # cpu
+    index_bp = {i:idx for idx, i in enumerate(index_fp.numpy().tolist())}
+    index_bp = [index_bp[i] for i in range(H * W)]
+    index_fp = paddle.to_tensor(index_fp, place=place)
+    index_bp = paddle.to_tensor(index_fp, dtype='int64', place=place)
+    return [index_fp, index_bp]
+
+class NpuRollWithIndexSelect():
+    def __init__(self):
+        self.index_dict = {}
+
+    def __call__(self, x, shifts, axis):
+        assert x.dim() == 4
+        assert len(shifts) == 2
+        assert len(axis) == 2
+        N, H, W, C = x.shape
+        key = (H, W, shifts, axis)
+        if key not in self.index_dict:
+            self.index_dict[key] = get_roll_index(H, W, shifts, x.place)
+        index_fp, index_bp = self.index_dict[key]
+        return roll_with_index_select(x, index_fp, index_bp)
+
+roll = NpuRollWithIndexSelect()
 
 class Mlp(nn.Layer):
     def __init__(self,
@@ -356,7 +400,7 @@ class SwinTransformerBlock(nn.Layer):
 
         # cyclic shift
         if self.shift_size > 0:
-            shifted_x = paddle.roll(
+            shifted_x = roll(
                 x, shifts=(-self.shift_size, -self.shift_size), axis=(1, 2))
         else:
             shifted_x = x
@@ -380,7 +424,7 @@ class SwinTransformerBlock(nn.Layer):
 
         # reverse cyclic shift
         if self.shift_size > 0:
-            x = paddle.roll(
+            x = roll(
                 shifted_x,
                 shifts=(self.shift_size, self.shift_size),
                 axis=(1, 2))