项目地址
SamOutV2 0.18B模型
- 采取 em参数共享在参数量减半的情况下将维度从1024 拉升到了1536
- sft 单论对话 loss 保持1.8
- 如果未来匹配state 推理代码性能不变的同时推理任意长度使用资源空间保持不变
模型代码
import torchclass MaxState(torch.nn.Module):def __init__(self, hidden_dim, heads):super(MaxState, self).__init__()assert hidden_dim % heads == 0, "Hidden size must be divisible by the number of heads."self.head_size = hidden_dim // headsself.head0 = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)self.head1 = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)self.head2 = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)self.head_num = headsself.hidden = hidden_dimdef forward(self, input_data, state=None):b, s, k, h = input_data.shape[0], input_data.shape[1], self.head_num, self.head_sizeout = self.head0(input_data)out1 = self.head1(input_data)out2 = self.head2(input_data)out = out.reshape([b, s, k, h]).permute([0, 2, 1, 3])out1 = out1.reshape([b, s, k, h]).permute([0, 2, 1, 3])out = torch.cummax((out + out1) / h ** 0.5, 2)[0]out = out.permute([0, 2, 1, 3])out1 = out1.permute([0, 2, 1, 3])out = out.reshape([b, s, -1])out1 = out1.reshape([b, s, -1])out = (out + out2) * out + out1return out, stateclass FeedForward(torch.nn.Module):def __init__(self, hidden_size):super(FeedForward, self).__init__()self.ffn1 = torch.nn.Linear(hidden_size, hidden_size // 2)self.ffn2 = torch.nn.Linear(hidden_size // 2, hidden_size)self.gate = torch.nn.Linear(hidden_size, hidden_size // 2)self.relu = torch.nn.ReLU()self.gr=torch.nn.Dropout(0.1)def forward(self, x):x1 = self.ffn1(x)x2 = self.relu(self.gate(x))xx = x1 * x2x = self.gr(self.ffn2(xx))return xclass DecoderLayer(torch.nn.Module):def __init__(self, hidden_size, num_heads):super(DecoderLayer, self).__init__()self.self_attention = MaxState(hidden_size, num_heads)self.ffn = FeedForward(hidden_size)self.layer_norm = torch.nn.LayerNorm(hidden_size)self.alpha = torch.nn.Parameter(torch.tensor(0.5))def forward(self, x, state=None, ):x1, state = self.self_attention(x, state)x = self.layer_norm(self.alpha * self.ffn(x1) + (1 - self.alpha) * x)return x, stateclass SamOut(torch.nn.Module):def __init__(self, voc_size, hidden_size, num_heads, num_layers):super(SamOut, self).__init__()self.em = torch.nn.Embedding(voc_size, hidden_size, padding_idx=3)self.decoder_layers = torch.nn.ModuleList([DecoderLayer(hidden_size, num_heads) for _ in range(num_layers)])self.head = FeedForward(hidden_size)def state_forward(self, state, x):if state is None:state = [None] * len(self.decoder_layers)i = 0for ii, decoder_layer in enumerate(self.decoder_layers):x1, state[i] = decoder_layer(x, state[i])x = x1 + xi += 1return x, statedef forward(self, x, state=None):x = self.em(x)x, state = self.state_forward(state, x)em = self.head(self.em.weight) / x.shape[-1]return x @ em.permute([1, 0]), stateif __name__ == '__main__':net = SamOut(235, 256, 16, 4)net(torch.randint(0, 200, [2, 8 * 13]))代码解析
这段代码定义了一个名为 SamOut 的神经网络模型，它包含了一些自定义的组件。这些组件共同工作以实现一个解码器（decoder），可以用于自然语言处理任务中的文本生成或翻译等。下面是对代码中各个类和方法的解释：
MaxState 类
MaxState 是一个多头机制的模块，它接收输入数据，并通过三个线性变换层 (head0, head1, head2) 来处理数据。这个模块的主要功能是计算累积最大值（cumulative maximum）并且在计算过程中考虑了缩放因子（h ** 0.5）。每个头的输出尺寸是 hidden_dim / heads，即 head_size。
FeedForward 类
FeedForward 是一个前馈神经网络（feed-forward network），它包含了两个线性变换层 (ffn1 和 ffn2) 和一个门控机制 (gate)。它还使用了 ReLU 激活函数和 dropout 层来防止过拟合。这个模块通常用作 Transformer 架构中的一部分，用来增加模型的表达能力。
DecoderLayer 类
DecoderLayer 包含了一个 MaxState 实例作为自注意力机制的一部分，以及一个 FeedForward 网络。此外，它还实现了残差连接（residual connection）和层归一化（layer normalization）。alpha 参数用于调整来自 self_attention 和 ffn 的贡献比例。
SamOut 类
SamOut 是整个解码器模型的主体，它整合了所有上述组件。它接受词汇表大小 (voc_size)、隐藏层大小 (hidden_size)、多头数 (num_heads) 和层数 (num_layers) 作为参数。它有一个嵌入层 (em) 来将输入token映射到向量空间，并且有一系列的 DecoderLayer 实例组成的模块列表 (decoder_layers)。最后，它还有一个 FeedForward 层 (head) 用于最终的输出转换。
Forward 方法
forward方法首先将输入token序列通过嵌入层转换为向量。- 然后，这些向量被传递给由多个
DecoderLayer组成的堆栈，每一层都会更新状态信息。 - 最终，通过
head层进行线性变换，并与嵌入矩阵相乘得到logits，这是为了预测下一个token的概率分布。
主程序部分
在主程序部分，创建了一个 SamOut 实例 net，并随机生成了一些整数作为输入来测试模型。这里的 torch.randint(0, 200, [2, 8 * 13]) 生成了一个形状为 [2, 104] 的张量，其中的元素是从 0 到 200 的随机整数，代表batch size为2的一批输入序列，每个序列长度为104。
请注意，这里展示的模型结构是高度定制化的，并不是标准的Transformer架构。特别是 MaxState 使用的是累积最大值而不是传统的多头自注意力机制。如果你打算基于此代码构建实际应用，请确保理解其工作机制，并根据你的具体需求进行必要的调整和优化。