වේගයෙන් සිරුරේ බර පද්ධති

9fromtypingimportOptional1011importtorch12fromtorchimportnn1314fromlabml\_helpers.moduleimportModule15fromlabml\_nn.transformers.fast\_weightsimportDPFP16fromlabml\_nn.transformers.feed\_forwardimportFeedForward17fromlabml\_nn.transformers.mhaimportPrepareForMultiHeadAttention18fromlabml\_nn.utilsimportclone\_module\_list

21classFastWeightsAttention(Module):

22def\_\_init\_\_(self,heads:int,d\_model:int,dropout\_prob:float,phi:DPFP):23super().\_\_init\_\_()

26self.d\_k=d\_model//heads

28self.heads=heads

31self.query=PrepareForMultiHeadAttention(d\_model,heads,self.d\_k,bias=False)

33self.key=PrepareForMultiHeadAttention(d\_model,heads,self.d\_k,bias=False)34self.value=PrepareForMultiHeadAttention(d\_model,heads,self.d\_k,bias=False)3536self.gate=nn.Sequential(PrepareForMultiHeadAttention(d\_model,heads,1,bias=False),37nn.Sigmoid())3839self.phi=phi

42self.output=nn.Linear(d\_model,d\_model)

44self.dropout=nn.Dropout(dropout\_prob)

46defforward(self,x:torch.Tensor,weights:Optional[torch.Tensor]):47query=self.phi(self.query(x))48key=self.phi(self.key(x))49value=self.value(x)5051ifweightsisNone:52weights=key.new\_zeros((key.shape[0],key.shape[1],value.shape[2],key.shape[2]))5354value\_existing=torch.einsum('bhvk,bhk-\>bhv',weights,key)5556beta=self.gate(x)5758weights=weights+torch.einsum('bhv,bhk-\>bhvk',beta\*(value-value\_existing),key)5960x=torch.einsum('bhvk,bhk-\>bhv',weights,query)

63x=x.reshape(x.shape[0],-1)

66returnself.output(x),weights

69classFastWeightsAttentionTransformerLayer(Module):

70def\_\_init\_\_(self,\*,71d\_model:int,72attn:FastWeightsAttention,73feed\_forward:FeedForward,74dropout\_prob:float):75super().\_\_init\_\_()

77self.size=d\_model

79self.attn=attn80self.feed\_forward=feed\_forward81self.dropout=nn.Dropout(dropout\_prob)

84self.norm\_self\_attn=nn.LayerNorm([d\_model])85self.norm\_ff=nn.LayerNorm([d\_model])

87defforward(self,x:torch.Tensor,weights:Optional[torch.Tensor]):88attn,weights=self.attn(x,weights)

90x=x+self.dropout(attn)

93z=self.norm\_ff(x)

95ff=self.feed\_forward(z)

97x=x+self.dropout(ff)

100returnx,weights

103classFastWeightsAttentionTransformer(Module):

104def\_\_init\_\_(self,layer:FastWeightsAttentionTransformerLayer,n\_layers:int):105super().\_\_init\_\_()

107self.layers=clone\_module\_list(layer,n\_layers)

109self.norm=nn.LayerNorm([layer.size])

111defforward(self,x\_seq:torch.Tensor):

113x\_seq=torch.unbind(x\_seq,dim=0)

115res=[]

117weights=[Nonefor\_inrange(len(self.layers))]118119forxinx\_seq:

121fori,layerinenumerate(self.layers):

123x,weights[i]=layer(x,weights[i])124125res.append(x)

128res=torch.stack(res)

130returnself.norm(res)