attention.py
CFM 上游 AttentionNet
原文件 · 324 行 · 14524 字节
third_party/heteromrta/attention.py
SHA-256: 7c214f8c51aee3df8acf3cd033f5229c88b8c1a84d9d76c91f38bd144e6f54ec
1import torch2import torch.nn as nn3import math4import numpy as np5from torch.nn.utils.rnn import pad_sequence6from torch.cuda.amp.autocast_mode import autocast7from parameters import *8910def get_attn_pad_mask(seq_q, seq_k):11 """生成 padding mask；整行都是零的输入不会参与注意力计算。"""12 batch_size, len_q = seq_q.sum(dim=2).size()13 batch_size, len_k = seq_k.sum(dim=2).size()14 # eq(zero) is PAD token15 pad_attn_mask_k = seq_q.eq(0).all(2).data.eq(1).unsqueeze(1) # batch_size x 1 x len_q, one is masking16 pad_attn_mask_q = seq_k.eq(0).all(2).data.eq(1).unsqueeze(1) # batch_size x 1 x len_k, one is masking17 pad_attn_mask_k = pad_attn_mask_k.expand(batch_size, len_k, len_q).permute(0, 2, 1)18 pad_attn_mask_q = pad_attn_mask_q.expand(batch_size, len_q, len_k)19 return ~torch.logical_and(~pad_attn_mask_k, ~pad_attn_mask_q) # batch_size x len_q x len_k202122def get_attn_subsequent_mask(seq):23 """生成按顺序解码时使用的下三角 mask。"""24 attn_shape = [seq.size(0), seq.size(1), seq.size(1)]25 subsequent_mask = np.logical_not(np.triu(np.ones(attn_shape), k=0)).astype(int)26 subsequent_mask = torch.from_numpy(subsequent_mask).byte()27 return subsequent_mask282930class SingleHeadAttention(nn.Module):31 """计算 query 对每个候选任务的概率，供最终任务选择使用。"""32 def __init__(self, embedding_dim):33 super(SingleHeadAttention, self).__init__()34 self.input_dim = embedding_dim35 self.embedding_dim = embedding_dim36 self.value_dim = embedding_dim37 self.key_dim = self.value_dim38 self.tanh_clipping = 1039 self.norm_factor = 1 / math.sqrt(self.key_dim)4041 self.w_query = nn.Parameter(torch.Tensor(self.input_dim, self.key_dim))42 self.w_key = nn.Parameter(torch.Tensor(self.input_dim, self.key_dim))4344 self.init_parameters()4546 def init_parameters(self):47 for param in self.parameters():48 stdv = 1. / math.sqrt(param.size(-1))49 param.data.uniform_(-stdv, stdv)5051 def forward(self, q, h=None, mask=None):52 """53 :param q: queries (batch_size, n_query, input_dim)54 :param h: data (batch_size, graph_size, input_dim)55 :param mask: mask (batch_size, n_query, graph_size) or viewable as that (i.e. can be 2 dim if n_query == 1)56 Mask should contain 1 if attention is not possible (i.e. mask is negative adjacency)57 :return:58 """59 if h is None:60 # 未单独提供 h 时，对 q 自身计算注意力。61 h = q6263 batch_size, target_size, input_dim = h.size()64 n_query = q.size(1) # n_query = target_size in tsp6566 h_flat = h.reshape(-1, input_dim) # (batch_size*graph_size)*input_dim67 q_flat = q.reshape(-1, input_dim) # (batch_size*n_query)*input_dim6869 shape_k = (batch_size, target_size, -1)70 shape_q = (batch_size, n_query, -1)7172 Q = torch.matmul(q_flat, self.w_query).view(shape_q) # batch_size*n_query*key_dim73 K = torch.matmul(h_flat, self.w_key).view(shape_k) # batch_size*targets_size*key_dim7475 U = self.norm_factor * torch.matmul(Q, K.transpose(1, 2)) # batch_size*n_query*targets_size76 # tanh 限制最终任务选择分数的范围。77 U = self.tanh_clipping * torch.tanh(U)7879 if mask is not None:80 mask = mask.view(batch_size, -1, target_size).expand_as(U) # copy for n_heads times81 U[mask.bool()] = -1e882 attention = torch.softmax(U, dim=-1) # batch_size*n_query*targets_size83 logp_list = torch.log_softmax(U, dim=-1) # batch_size*n_query*targets_size8485 probs = attention8687 return probs, logp_list888990class MultiHeadAttention(nn.Module):91 """用多个 head 汇总输入序列，再拼回原来的特征维度。"""92 def __init__(self, embedding_dim, n_heads=8):93 super(MultiHeadAttention, self).__init__()94 self.n_heads = n_heads95 self.input_dim = embedding_dim96 self.embedding_dim = embedding_dim97 self.value_dim = self.embedding_dim // self.n_heads98 self.key_dim = self.value_dim99 self.norm_factor = 1 / math.sqrt(self.key_dim)100101 self.w_query = nn.Parameter(torch.Tensor(self.n_heads, self.input_dim, self.key_dim))102 self.w_key = nn.Parameter(torch.Tensor(self.n_heads, self.input_dim, self.key_dim))103 self.w_value = nn.Parameter(torch.Tensor(self.n_heads, self.input_dim, self.value_dim))104 self.w_out = nn.Parameter(torch.Tensor(self.n_heads, self.value_dim, self.embedding_dim))105106 self.init_parameters()107108 def init_parameters(self):109 for param in self.parameters():110 stdv = 1. / math.sqrt(param.size(-1))111 param.data.uniform_(-stdv, stdv)112113 def forward(self, q, h=None, mask=None):114 """115 :param q: queries (batch_size, n_query, input_dim)116 :param h: data (batch_size, graph_size, input_dim)117 :param mask: mask (batch_size, n_query, graph_size) or viewable as that (i.e. can be 2 dim if n_query == 1)118 Mask should contain 1 if attention is not possible (i.e. mask is negative adjacency)119 :return:120 """121 if h is None:122 h = q123124 batch_size, target_size, input_dim = h.size()125 n_query = q.size(1) # n_query = target_size in tsp126127 h_flat = h.contiguous().view(-1, input_dim) # (batch_size*graph_size)*input_dim128 q_flat = q.contiguous().view(-1, input_dim) # (batch_size*n_query)*input_dim129 shape_v = (self.n_heads, batch_size, target_size, -1)130 shape_k = (self.n_heads, batch_size, target_size, -1)131 shape_q = (self.n_heads, batch_size, n_query, -1)132133 # 三个投影先按 head 拆分，再计算每个 query 对所有输入行的权重。134 Q = torch.matmul(q_flat, self.w_query).view(shape_q) # n_heads*batch_size*n_query*key_dim135 K = torch.matmul(h_flat, self.w_key).view(shape_k) # n_heads*batch_size*targets_size*key_dim136 V = torch.matmul(h_flat, self.w_value).view(shape_v) # n_heads*batch_size*targets_size*value_dim137138 U = self.norm_factor * torch.matmul(Q, K.transpose(2, 3)) # n_heads*batch_size*n_query*targets_size139140 if mask is not None:141 mask = mask.view(1, batch_size, -1, target_size).expand_as(U) # copy for n_heads times142 # U[mask.bool()] = -np.inf143 U[mask.bool()] = -np.inf144 attention = torch.softmax(U, dim=-1) # n_heads*batch_size*n_query*targets_size145146 if mask is not None:147 attnc = attention.clone()148 attnc[mask.bool()] = 0149 attention = attnc150 # print(attention)151152 heads = torch.matmul(attention, V) # n_heads*batch_size*n_query*value_dim153154 # 将所有 head 拼接后投影回 embedding_dim。155 out = torch.mm(156 heads.permute(1, 2, 0, 3).reshape(-1, self.n_heads * self.value_dim),157 # batch_size*n_query*n_heads*value_dim158 self.w_out.view(-1, self.embedding_dim)159 # n_heads*value_dim*embedding_dim160 ).view(batch_size, n_query, self.embedding_dim)161162 return out # batch_size*n_query*embedding_dim163164165class GateFFNDense(nn.Module):166 """用一条 sigmoid 分支控制另一条线性分支的输出大小。"""167 def __init__(self, model_dim, hidden_unit=512):168 super(GateFFNDense, self).__init__()169 self.W = nn.Linear(model_dim, hidden_unit, bias=False)170 self.V = nn.Linear(model_dim, hidden_unit, bias=False)171 self.W2 = nn.Linear(hidden_unit, model_dim, bias=False)172 self.act = nn.Sigmoid()173174 def forward(self, hidden_states):175 hidden_act = self.act(self.W(hidden_states))176 hidden_linear = self.V(hidden_states)177 hidden_states = hidden_act * hidden_linear178 hidden_states = self.W2(hidden_states)179 return hidden_states180181182class GateFFNLayer(nn.Module):183 def __init__(self, model_dim):184 super(GateFFNLayer, self).__init__()185 self.DenseReluDense = GateFFNDense(model_dim)186 self.layer_norm = Normalization(model_dim)187188 def forward(self, hidden_states):189 forwarded_states = self.layer_norm(hidden_states)190 forwarded_states = self.DenseReluDense(forwarded_states)191 return forwarded_states192193194class Normalization(nn.Module):195 def __init__(self, embedding_dim):196 super(Normalization, self).__init__()197 self.normalizer = nn.LayerNorm(embedding_dim)198199 def forward(self, input):200 return self.normalizer(input.view(-1, input.size(-1))).view(*input.size())201202203class EncoderLayer(nn.Module):204 """依次执行自注意力和前馈网络，并在两步后加入残差。"""205 def __init__(self, embedding_dim, n_head):206 super(EncoderLayer, self).__init__()207 self.multiHeadAttention = MultiHeadAttention(embedding_dim, n_head)208 self.normalization1 = Normalization(embedding_dim)209 self.feedForward = GateFFNLayer(embedding_dim)210211 def forward(self, src, mask=None):212 h0 = src213 h = self.normalization1(src)214 h = self.multiHeadAttention(q=h, mask=mask)215 h = h + h0216 h1 = h217 h = self.feedForward(h)218 h = h + h1219 return h220221222class DecoderLayer(nn.Module):223 """用当前智能体表示读取已经编码的任务表示。"""224 def __init__(self, embedding_dim, n_head):225 super(DecoderLayer, self).__init__()226 self.dec_self_attn = MultiHeadAttention(embedding_dim, n_head)227 self.multiHeadAttention = MultiHeadAttention(embedding_dim, n_head)228 self.feedForward = GateFFNLayer(embedding_dim)229 self.normalization1 = Normalization(embedding_dim)230 self.normalization2 = Normalization(embedding_dim)231232 def forward(self, tgt, memory, dec_self_attn_mask, dec_enc_attn_mask):233 h0 = tgt234 tgt = self.normalization1(tgt)235 memory = self.normalization2(memory)236 h = self.multiHeadAttention(q=tgt, h=memory, mask=dec_enc_attn_mask)237 h = h + h0238 h1 = h239 h = self.feedForward(h)240 h = h + h1241 return h242243244class Encoder(nn.Module):245 def __init__(self, embedding_dim=128, n_head=4, n_layer=2):246 super(Encoder, self).__init__()247 self.layers = nn.ModuleList(EncoderLayer(embedding_dim, n_head) for i in range(n_layer))248249 def forward(self, src, mask=None):250 # 每层都保持相同的序列长度和 embedding 宽度。251 for layer in self.layers:252 src = layer(src, mask)253 return src254255256class Decoder(nn.Module):257 def __init__(self, embedding_dim=128, n_head=4, n_layer=2):258 super(Decoder, self).__init__()259 self.layers = nn.ModuleList([DecoderLayer(embedding_dim, n_head) for i in range(n_layer)])260261 def forward(self, tgt, memory, dec_self_attn_mask=None, dec_enc_attn_mask=None):262 for layer in self.layers:263 tgt = layer(tgt, memory, dec_self_attn_mask, dec_enc_attn_mask)264 return tgt265266267class AttentionNet(nn.Module):268 def __init__(self, agent_input_dim, task_input_dim, embedding_dim):269 super(AttentionNet, self).__init__()270 self.agent_embedding = nn.Linear(agent_input_dim, embedding_dim)271 self.task_embedding = nn.Linear(task_input_dim, embedding_dim) # layer for input information272 self.fusion = nn.Linear(embedding_dim * 3, embedding_dim)273274 self.taskEncoder = Encoder(embedding_dim=embedding_dim, n_head=8, n_layer=1)275 self.crossDecoder1 = Decoder(embedding_dim=embedding_dim, n_head=8, n_layer=2)276 self.crossDecoder2 = Decoder(embedding_dim=embedding_dim, n_head=8, n_layer=2)277 self.agentEncoder = Encoder(embedding_dim=embedding_dim, n_head=8, n_layer=1)278 self.globalDecoder = Decoder(embedding_dim=embedding_dim, n_head=8, n_layer=2)279 self.pointer = SingleHeadAttention(embedding_dim)280 # self.LSTM = nn.LSTM(embedding_dim, embedding_dim, batch_first=True)281282 def encoding_tasks(self, task_inputs, mask=None):283 task_embedding = self.task_embedding(task_inputs)284 task_encoding = self.taskEncoder(task_embedding, mask)285 embedding_dim = task_encoding.size(-1)286 mean_mask = mask[:,0,:].unsqueeze(2).repeat(1, 1, embedding_dim)287 compressed_task = torch.where(mean_mask, torch.nan, task_embedding)288 aggregated_tasks = torch.nanmean(compressed_task, dim=1).unsqueeze(1)289 return aggregated_tasks, task_encoding290291 def encoding_agents(self, agents_inputs, mask=None):292 agents_embedding = self.agent_embedding(agents_inputs)293 agents_encoding = self.agentEncoder(agents_embedding, mask)294 embedding_dim = agents_encoding.size(-1)295 mean_mask = mask[:,0,:].unsqueeze(2).repeat(1, 1, embedding_dim)296 compressed_task = torch.where(mean_mask, torch.nan, agents_embedding)297 aggregated_agents = torch.nanmean(compressed_task, dim=1).unsqueeze(1)298 return aggregated_agents, agents_encoding299300 def forward(self, tasks, agents, global_mask, index):301 task_mask = get_attn_pad_mask(tasks, tasks)302 agent_mask = get_attn_pad_mask(agents, agents)303 task_agent_mask = get_attn_pad_mask(tasks, agents)304 agent_task_mask = get_attn_pad_mask(agents, tasks)305 aggregated_task, task_encoding = self.encoding_tasks(tasks, mask=task_mask)306 aggregated_agents, agents_encoding = self.encoding_agents(agents, mask=agent_mask)307 task_agent_feature = self.crossDecoder1(task_encoding, agents_encoding, None, task_agent_mask)308 agent_task_feature = self.crossDecoder2(agents_encoding, task_encoding, None, agent_task_mask)309 current_state1 = torch.gather(agent_task_feature, 1, index.repeat(1, 1, agent_task_feature.size(2)))310 current_state = self.fusion(torch.cat((current_state1, aggregated_task, aggregated_agents), dim=-1))311 current_state_prime = self.globalDecoder(current_state, task_agent_feature, None, global_mask)312 probs, logps = self.pointer(current_state_prime, task_agent_feature, mask=global_mask)313 logps = logps.squeeze(1)314 probs = probs.squeeze(1)315 return probs, logps316317318def padding_inputs(inputs):319 seq = pad_sequence(inputs, batch_first=False, padding_value=1)320 seq = seq.permute(2, 1, 0)321 mask = torch.zeros_like(seq, dtype=torch.int64)322 ones = torch.ones_like(seq, dtype=torch.int64)323 mask = torch.where(seq != 1, mask, ones)324 return seq, mask