在Transformer模型中,Feed Forward Neural Network (FFNN) 是由两个线性层和一个非线性激活函数(通常是ReLU)组成的。以下是使用PyTorch实现Transformer中Feed Forward部分的示例代码:
python复制代码
import torch
import torch.nn as nn
import torch.nn.functional as F
class FeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super(FeedForward, self).__init__()
# 两个线性层:第一层将输入维度d_model映射到d_ff,第二层将d_ff映射回d_model
self.linear1 = nn.Linear(d_model, d_ff)
self.dropout = nn.Dropout(dropout)
self.linear2 = nn.Linear(d_ff, d_model)
def forward(self, x):
# 第一个线性层
x = self.linear1(x)
# 应用ReLU激活函数
x = F.relu(x)
# 应用dropout
x = self.dropout(x)
# 第二个线性层
x = self.linear2