samout 结构再优化 收敛速度再加速

代码

1. import torch
2. import numpy as np
5. class MaxState(torch.nn.Module):
6.     def __init__(self, hidden_dim, heads, win):
7.         super(MaxState, self).__init__()
9.         assert hidden_dim % heads == 0, "Hidden size must be divisible by the number of heads."
12.         self.head_size = hidden_dim // heads
13.         self.head = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
14.         self.state = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
15.         self.head_num = heads
16.         self.win = win
17.         self.hidden = hidden_dim
18.         self.mask = torch.triu(torch.ones([win, win])).to("cuda")
19.         self.layer_nor = torch.nn.LayerNorm(hidden_dim)
21.     def forward(self, input_data, state=None):
22.         # self.head.to("cuda")
23.         b, s, k, h, w = input_data.shape[0], input_data.shape[1], self.head_num, self.head_size, self.win
25.         window = torch.ones([1, w]).to("cuda")
28.         out = self.head(input_data)
31.         out = out.unsqueeze(-1) @ window
34.         out = out.permute([0, 2, 1, 3])
37.         one_list = []
38.         if state is None:
39.             state = torch.ones([out.shape[0], out.shape[1], 1, 1]) * float("-inf")
40.             state = state.to("cuda")
41.         for i in range(0, s, w):
43.             state.reshape([state.shape[0], -1])
44.             j = w + i
45.             one = out[:, :, i:j]
46.             _, _, r, c = one.shape
47.             if r != self.win:
49.                 one = torch.where(self.mask[:r, :] == 1, one, torch.Tensor([-float('inf')]).to("cuda"))
51.             else:
52.                 one = torch.where(self.mask == 1, one, torch.Tensor([-float('inf')]).to("cuda"))
54.             if i == 0:
56.                 one = torch.concat([one, state @ window], axis=2)
57.                 state, _ = torch.max(one, axis=2, keepdim=True)
60.             else:
62.                 state1, _ = torch.max(one, axis=2, keepdim=True)
64.                 # state = torch.sin(self.state(state1.reshape([state1.shape[0], -1]))*state.reshape([state.shape[0], -1]))
65.                 state1 = self.state(state1.permute([0, 3, 1, 2]).reshape([state1.shape[0], -1, state1.shape[1]]))
66.                 state = state1.permute([0, 2, 1]).unsqueeze(-2) + state
67.                 # state = state.reshape(state1.shape)
69.                 one = torch.concat([one, state], axis=2)
70.                 state, _ = torch.max(one, axis=2, keepdim=True)
72.             one = state.reshape([b, k, h, w])
74.             state = state[..., -1:]
75.             if r != self.win:
76.                 one = one[..., :r]
78.             one = one.permute([0, 3, 1, 2])
79.             one_list.append(one)
81.         out = torch.concat(one_list, 1)
83.         out = out.reshape([b, s, -1])
85.         return out, state
89. class FeedForward(torch.nn.Module):
90.     def __init__(self, hidden_size):
91.         super(FeedForward, self).__init__()
93.         self.ffn1 = torch.nn.Linear(hidden_size, hidden_size * 2)
94.         self.ffn2 = torch.nn.Linear(hidden_size * 2, hidden_size)
95.         self.gate = torch.nn.Linear(hidden_size, hidden_size * 2)
96.         self.relu = torch.nn.ReLU()
98.     def forward(self, x):
99.         x1 = self.ffn1(x)
101.         x2 = self.relu(self.gate(x))
103.         x = x1 * x2
105.         x = self.ffn2(x)
106.         return x
109. class DecoderLayer(torch.nn.Module):
110.     def __init__(self, hidden_size, num_heads):
111.         super(DecoderLayer, self).__init__()
112.         # self.self_attention = MaskMultiHeadAttention(hidden_size, num_heads)
113.         self.self_attention = MaxState(hidden_size, num_heads, 8)
114.         self.ffn = FeedForward(hidden_size)
115.         self.layer_norm = torch.nn.LayerNorm(hidden_size)
117.     def forward(self, x, state=None, seq_len=None):
118.         x1, state = self.self_attention(x, state)
120.         x = self.layer_norm(self.ffn(x1) + x)  # Feed-Forward with residual connection
122.         return x, state
125. class SamOut(torch.nn.Module):
126.     def __init__(self, voc_size, hidden_size, num_heads, num_layers):
127.         super(SamOut, self).__init__()
128.         self.em = torch.nn.Embedding(voc_size, hidden_size, padding_idx=3)
129.         self.pos = torch.nn.Embedding(1024, hidden_size)
131.         self.decoder_layers = torch.nn.ModuleList([DecoderLayer(hidden_size, num_heads) for _ in range(num_layers)])
132.         self.head = torch.nn.Linear(hidden_size, voc_size)
133.         self.head_state = torch.nn.Linear(hidden_size, num_layers)
134.         self.layer_nor=torch.nn.LayerNorm(hidden_size)
135.         self.down=torch.nn.ModuleList([torch.nn.Linear(2*hidden_size,hidden_size) for _ in range(num_layers)])
137.     def forward(self, x, state=None, seq_len=None):
138.         x = self.em(x)
139.         if x.shape[1] >= 1024:
140.             pos = self.pos(torch.range(0, x.shape[1]-1).long() // 1024).unsqueeze(0)
141.             pos = self.pos(torch.range(0, x.shape[1]-1).long() % 1024).unsqueeze(0) + pos
143.         else:
144.             pos = self.pos(torch.range(0, x.shape[1]-1).long().to("cuda")).unsqueeze(0)
146.         if state is None:
147.             state = [None] * len(self.decoder_layers)
149.         i = 0
151.         for decoder_layer in self.decoder_layers:
152.             x1, state[i] = decoder_layer(self.down[i](torch.concat([torch.zeros([x.shape[0],1,1]).to("cuda")+pos , x],-1)), state[i])
153.             x = x1 + x
155.             i += 1
156.         state_data = self.head_state((torch.concat(state, -1).squeeze(-2)).permute([0, 2, 1]))
157.         return self.head(x), state, state_data
160. if __name__ == '__main__':
161.     net = SamOut(235, 256, 16, 4)
162.     net(torch.randint(0, 200, [2, 3000]))

复制代码

解释

这段代码定义了一个基于PyTorch的神经网络模型，该模型包罗自定义的解码器层和输出层，用于处置惩罚序列数据。下面是代码的逐行解析：

1. import torch
2. import numpy as np

复制代码

• 导入PyTorch库和NumPy库。
  

1. class MaxState(torch.nn.Module):
2.     def __init__(self, hidden_dim, heads, win):
3.         super(MaxState, self).__init__()

复制代码

• 定义一个名为MaxState的类，继续自torch.nn.Module。这是自定义的一个模块，用于处置惩罚状态的最大化。
  

1.         assert hidden_dim % heads == 0, "Hidden size must be divisible by the number of heads."

复制代码

• 断言查抄，确保隐蔽层的维度可以被注意力头的数目整除。
  

1.         self.head_size = hidden_dim // heads
2.         self.head = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
3.         self.state = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
4.         self.head_num = heads
5.         self.win = win
6.         self.hidden = hidden_dim
7.         self.mask = torch.triu(torch.ones([win, win])).to("cuda")
8.         self.layer_nor = torch.nn.LayerNorm(hidden_dim)

复制代码

• 初始化类成员变量，包罗线性层、注意力头数目、窗口巨细、掩码和层归一化。
  

1.     def forward(self, input_data, state=None):
2.         # self.head.to("cuda")
3.         b, s, k, h, w = input_data.shape[0], input_data.shape[1], self.head_num, self.head_size, self.win

复制代码

• 前向传播方法，获取输入数据的形状参数。
  

1.         window = torch.ones([1, w]).to("cuda")

复制代码

• 创建一个窗口张量，并将其移动到CUDA装备。
  

1.         out = self.head(input_data)

复制代码

• 应用线性层到输入数据。
  

1.         out = out.unsqueeze(-1) @ window

复制代码

• 扩展维度并举行矩阵乘法。
  

1.         out = out.permute([0, 2, 1, 3])

复制代码

• 调解输出的维度顺序。
  

1.         one_list = []
2.         if state is None:
3.             state = torch.ones([out.shape[0], out.shape[1], 1, 1]) * float("-inf")
4.             state = state.to("cuda")

复制代码

• 初始化状态张量，如果状态为None，则创建一个初始状态。
  

1.         for i in range(0, s, w):
2.             # ... (省略中间代码)

复制代码

• 循环处置惩罚每个窗口巨细的数据块。
  

1.         return out, state

复制代码

• 返回处置惩罚后的输出和状态。
  接下来是FeedForward、DecoderLayer和SamOut类的定义，这些类分别实现了前馈网络、解码器层和整个模型的输出部分。代码结构与MaxState类雷同，包罗了初始化和前向传播方法。
  

1. if __name__ == '__main__':
2.     net = SamOut(235, 256, 16, 4)
3.     net(torch.randint(0, 200, [2, 3000]))

复制代码

• 如果当前脚本作为主程序运行，创建一个SamOut模型实例，并使用随机整数张量作为输入举行测试。
  

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