ixiaohuihuihui / gcc-net Goto Github PK

        # About 2x faster than original impl
        Wh, Ww = self.window_size
        rel_index_coords = self.double_step_seq(2 * Ww - 1, Wh, 1, Ww)
        rel_position_index = rel_index_coords + rel_index_coords.T
        rel_position_index = rel_position_index.flip(1).contiguous()
        self.register_buffer('relative_position_index', rel_position_index)
        @staticmethod
        def double_step_seq(step1, len1, step2, len2):
             seq1 = torch.arange(0, step1 * len1, step1)
             seq2 = torch.arange(0, step2 * len2, step2)
             return (seq1[:, None] + seq2[None, :]).reshape(1, -1)

        B, N, C = x.shape
        # self.qkv生成qkv output of self.qkv: [5640, 49, 288]
        qkv_x = self.qkv(x).reshape(B, N, 3, self.num_heads,
                                  C // self.num_heads).permute(2, 0, 3, 1, 4)
        qkv_y = self.qkv(y).reshape(B, N, 3, self.num_heads,
                                  C // self.num_heads).permute(2, 0, 3, 1, 4)
        # qkv.shape = [3, 5640, 3, 49, 32]
        # make torchscript happy (cannot use tensor as tuple)
        qx, kx, vx = qkv_x[0], qkv_x[1], qkv_x[2]
        qy, ky, vy = qkv_y[0], qkv_y[1], qkv_y[2]

        qx = qx * self.scale
        # qy = qy * self.scale
        attn = (qx @ ky.transpose(-2, -1))

        # 定义一个相对位置的表:
        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index.view(-1)].view(
                self.window_size[0] * self.window_size[1],
                self.window_size[0] * self.window_size[1],
                -1)  # Wh*Ww,Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B // nW, nW, self.num_heads, N,
                             N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
        attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ vx).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)

        y = (attn @ vy).transpose(1, 2).reshape(B, N, C)
        y = self.proj(y)
        y = self.proj_drop(y)
        return x, y