调整残差块和分类头的 dropout 概率,并新增残差模块到 MoE 模型
This commit is contained in:
parent
e91f823d65
commit
fd913748ca
|
|
@ -58,7 +58,7 @@ EXPORT_HIDE_DIM = {
|
||||||
|
|
||||||
# ---------------------------- 残差块 ----------------------------
|
# ---------------------------- 残差块 ----------------------------
|
||||||
class ResidualBlock(nn.Module):
|
class ResidualBlock(nn.Module):
|
||||||
def __init__(self, dim, dropout_prob=0.1):
|
def __init__(self, dim, dropout_prob=0.0):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.linear1 = nn.Linear(dim, dim)
|
self.linear1 = nn.Linear(dim, dim)
|
||||||
self.ln1 = nn.LayerNorm(dim)
|
self.ln1 = nn.LayerNorm(dim)
|
||||||
|
|
@ -73,7 +73,7 @@ class ResidualBlock(nn.Module):
|
||||||
x = self.ln1(x)
|
x = self.ln1(x)
|
||||||
x = self.linear2(x)
|
x = self.linear2(x)
|
||||||
x = self.ln2(x)
|
x = self.ln2(x)
|
||||||
x = self.dropout(x) # 残差前加 Dropout(符合原描述)
|
x = self.dropout(x)
|
||||||
x = x + residual
|
x = x + residual
|
||||||
return self.relu(x)
|
return self.relu(x)
|
||||||
|
|
||||||
|
|
@ -86,7 +86,7 @@ class Expert(nn.Module):
|
||||||
d_model=1024,
|
d_model=1024,
|
||||||
num_resblocks=4,
|
num_resblocks=4,
|
||||||
output_multiplier=2,
|
output_multiplier=2,
|
||||||
dropout_prob=0.1,
|
dropout_prob=0.0,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
input_dim : BERT 输出的 hidden_size(如 312/768)
|
input_dim : BERT 输出的 hidden_size(如 312/768)
|
||||||
|
|
@ -156,6 +156,8 @@ class MoEModel(nn.Module):
|
||||||
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=4)
|
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=4)
|
||||||
self.pooler = nn.AdaptiveAvgPool1d(1)
|
self.pooler = nn.AdaptiveAvgPool1d(1)
|
||||||
|
|
||||||
|
self.res_blocks = nn.ModuleList([ResidualBlock(self.hidden_size) for _ in range(4)])
|
||||||
|
|
||||||
self.total_experts = 20
|
self.total_experts = 20
|
||||||
self.experts = nn.ModuleList()
|
self.experts = nn.ModuleList()
|
||||||
|
|
||||||
|
|
@ -175,6 +177,7 @@ class MoEModel(nn.Module):
|
||||||
|
|
||||||
# 4. 分类头
|
# 4. 分类头
|
||||||
self.classifier = nn.Sequential(
|
self.classifier = nn.Sequential(
|
||||||
|
nn.Dropout(0.2),
|
||||||
nn.LayerNorm(self.output_multiplier * self.hidden_size),
|
nn.LayerNorm(self.output_multiplier * self.hidden_size),
|
||||||
nn.Linear(
|
nn.Linear(
|
||||||
self.output_multiplier * self.hidden_size,
|
self.output_multiplier * self.hidden_size,
|
||||||
|
|
@ -186,10 +189,8 @@ class MoEModel(nn.Module):
|
||||||
self.output_multiplier * self.hidden_size * 2,
|
self.output_multiplier * self.hidden_size * 2,
|
||||||
),
|
),
|
||||||
nn.ReLU(inplace=True),
|
nn.ReLU(inplace=True),
|
||||||
nn.Dropout(0.2),
|
|
||||||
nn.Linear(self.output_multiplier * self.hidden_size * 2, num_classes),
|
nn.Linear(self.output_multiplier * self.hidden_size * 2, num_classes),
|
||||||
)
|
)
|
||||||
# 可选:为领域专家和共享专家设置不同权重衰减(通过优化器实现,此处不处理)
|
|
||||||
|
|
||||||
def to(self, device):
|
def to(self, device):
|
||||||
"""重写 to 方法,记录设备"""
|
"""重写 to 方法,记录设备"""
|
||||||
|
|
@ -212,6 +213,9 @@ class MoEModel(nn.Module):
|
||||||
embeddings, src_key_padding_mask=padding_mask
|
embeddings, src_key_padding_mask=padding_mask
|
||||||
) # [B, S, H]
|
) # [B, S, H]
|
||||||
|
|
||||||
|
for block in self.res_blocks:
|
||||||
|
encoded = block(encoded)
|
||||||
|
|
||||||
# ----- 3. 池化量 -----
|
# ----- 3. 池化量 -----
|
||||||
pooled = self.pooler(encoded.transpose(1, 2)).squeeze(-1)
|
pooled = self.pooler(encoded.transpose(1, 2)).squeeze(-1)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue