liwenyun
/
CLIP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix inefficient attention computation

This commit is contained in:
Penn 2022-07-21 11:35:03 -07:00
parent 4d120f3ec3
commit 8c0b98adfc
1 changed files with 4 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
clip/model.py
View File

@ -70,7 +70,7 @@ class AttentionPool2d(nn.Module):
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC
x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
x, _ = F.multi_head_attention_forward( x, _ = F.multi_head_attention_forward(
query=x, key=x, value=x, query=x[:1], key=x, value=x,
embed_dim_to_check=x.shape[-1], embed_dim_to_check=x.shape[-1],
num_heads=self.num_heads, num_heads=self.num_heads,
q_proj_weight=self.q_proj.weight, q_proj_weight=self.q_proj.weight,
@ -88,8 +88,7 @@ class AttentionPool2d(nn.Module):
training=self.training, training=self.training,
need_weights=False need_weights=False
) )
return x.squeeze(0)
return x[0]
class ModifiedResNet(nn.Module): class ModifiedResNet(nn.Module):