大模型条记：pytorch实现MOE

0 导入库

1. import torch
2. import torch.nn as nn
3. import torch.nn.functional as F

复制代码

1  专家模型

1. #一个简单的专家模型，可以是任何神经网络架构
2. class Expert(nn.Module):
3.     def __init__(self, input_size, output_size):
4.         super(Expert, self).__init__()
5.         self.fc = nn.Linear(input_size, output_size)
7.    
8.     def forward(self, x):
9.         return self.fc(x)

复制代码

2 MOE

1. class MoE(nn.Module):
2.     def __init__(self, num_experts, input_size, output_size,topk):
3.         super(MoE, self).__init__()
4.         self.num_experts = num_experts
5.         self.topk=topk
6.         self.experts = nn.ModuleList([Expert(input_size, output_size) for _ in range(num_experts)])
7.         #创建多个专家
8.         self.gating_network = nn.Linear(input_size, num_experts)  
9.         #门控网络
11.     def forward(self, x):
12.         #假设x的维度是(batch,input_size)
13.                   
14.         gating_scores = self.gating_network(x)
15.         # 门控网络决定权重 （选择每一个专家的概率）
16.         #输出维度是(batch_size, num_experts)
18.         topk_gate_scores,topk_gate_index=gating_scores.topk(topk,-1)
19.         #选取topk个专家
20.         #(batch_size,topk)
22.         gating_scores_filtered=torch.full_like(gating_scores,fill_value=float("-inf"))
23.         gating_scores_filtered=gating_scores_filtered.scatter(-1,topk_gate_index,topk_gate_scores)
24.         gating_scores_filtered=F.softmax(gating_scores_filtered,dim=-1)
25.         ##创建一个全为负无穷的张量 zeros，并将 topk_gate_scores 的值插入到这个张量的对应位置
26.         #(batch_size,num_experts)
27.                     
28.         expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1)  
29.         # 专家网络输出
30.         #每个expert的输出维度是 (batch_size, output_size)
31.         #stack沿着第二个维度堆叠，之后expert_outputs的维度是(batch_size,num_experts,output_size)
32.                   
33.         moe_output = torch.bmm(gating_scores_filtered.unsqueeze(1), expert_outputs).squeeze(1)
34.         # 加权组合专家输出
35.         #gating_scores.unsqueeze(1)——>(batch_size, 1,num_experts)
36.         #torch.bmm(gating_scores.unsqueeze(1), expert_outputs)——>(batch_size,1,output_size)
37.         #moe_output——>(batch_size,output_size)
38.         return moe_output

复制代码

3  输入举例

2. input_size = 10  
3. # 输入特征是大小为10的向量
4. output_size = 5  
5. # 输出大小为5的向量
6. num_experts = 3  
7. # 3个专家
10. moe_model = MoE(num_experts, input_size, output_size)
11. # 初始化MOE模型
14. input_vector = torch.randn(1, input_size)
15. # 创建一个输入向量
18. output_vector = moe_model(input_vector)
19. # 前向传递
22. print(output_vector.shape,output_vector)
23. # 打印输出
24. '''
25. torch.Size([1, 5]) tensor([[ 2.7343e-04,  4.0966e-01, -3.6634e-01, -8.9064e-01,  4.0759e-01]],
26.        grad_fn=<SqueezeBackward1>)
27. '''

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。