调整残差块和分类头的 dropout 概率,并新增残差模块到 MoE 模型

This commit is contained in:
songsenand 2026-02-15 00:08:44 +08:00
parent e91f823d65
commit fd913748ca
1 changed files with 9 additions and 5 deletions

View File

@ -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)