赞
踩
本文简单实现了标准的Transformer,也是我平时使用的Transformer基础代码。
导入相关包
import torch
import torch.nn as nn
from einops import rearrange
eniops操作张量维度非常方便,参考这篇文章 einops:优雅地操作张量维度
class MultiHeadAttention(nn.Module): def __init__(self,dim,num_heads,dim_head): """ 输入:(b, n, dim) dim: 输入序列的向量维度 num_heads: 注意力头的个数 dim_head: 每个注意力头的维度 """ super(MultiHeadAttention, self).__init__() self.num_heads = num_heads self.scale = dim_head ** -0.5 # weight_dim : q,k,v的维度 weight_dim = num_heads * dim_head self.qkv = nn.Linear(dim, weight_dim*3) self.proj = nn.Linear(weight_dim,dim) def forward(self,x): qkv = self.qkv(x).chunk(3,dim=-1) q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.num_heads), qkv) dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale attn = torch.softmax(dots,dim=-1) out = torch.matmul(attn,v) out = rearrange(out,'b h n d -> b n (h d)') out = self.proj(out) return out
class FFN(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class PreNorm(nn.Module): def __init__(self, dim, fn): super().__init__() self.norm = nn.LayerNorm(dim) self.fn = fn def forward(self, x, **kwargs): return self.fn(self.norm(x), **kwargs) class Transformer(nn.Module): def __init__(self, dim, num_heads, depth, embed_dim, mlp_dim, dropout=0.1): super().__init__() self.layers = nn.ModuleList([]) dim_head = embed_dim // num_heads for _ in range(depth): self.layers.append(nn.ModuleList([ PreNorm(dim, MultiHeadAttention(dim, num_heads, dim_head)), PreNorm(dim, FFN(dim, mlp_dim, dropout)) ])) def forward(self, x): for attn, ff in self.layers: x = attn(x) + x x = ff(x) + x return x
这里的Transformer没有添加位置编码,实际使用时需自行引入位置信息。
位置编码示例
# 简单起见,使用可学习的位置编码
b, n, d = 3, 1000, 64
x = torch.randn(b, n, d) # 输入序列
# 位置编码
pos_embed = nn.Parameter(torch.zeros(1, n, d))
nn.init.trunc_normal_(pos_embed, std=0.02)
# 添加位置编码
x = x + pos_embed
测试
if __name__ == '__main__':
x = torch.randn(3, 1000, 64)
net = Transformer(64,4,2,128,256)
y = net(x)
print(y.shape)
参考文献
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。