Method forward

diff2flow/kl_autoencoder.py:201–228

(self, x)

        self.use_null_attention = use_null_attention

    def forward(self, x):
        h_ = x
        h_ = self.norm(h_)
        q = self.q(h_)
        k = self.k(h_)
        v = self.v(h_)

        # compute null attention by discarding q,k
        if self.use_null_attention:
            out = self.proj_out(v)
            return x+out

        # compute attention
        if NATTEN_IS_AVAILBLE and self.natten_kernel_size > -1:
            q, k, v = map(lambda x: rearrange(x, 'b c h w -> b 1 h w c'), (q, k, v))
            qk = natten.functional.natten2dqk(q, k, self.natten_kernel_size, 1)
            a = torch.softmax(qk, dim=-1)
            out = natten.functional.natten2dav(a, v, self.natten_kernel_size, 1)
            out = rearrange(out, 'b 1 h w c -> b c h w')
        
        else:
            b,c,h,w = q.shape
            q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v))
            out = nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
            out = rearrange(out, 'b (h w) c -> b c h w', h=h, w=w)

        out = self.proj_out(out)
        return x+out


def make_attn(in_channels, attn_type="vanilla", natten_kernel_size=-1, use_null_attention=False):