import torch import torch.nn as nn from model.PDG2SEQ.PDG2SeqCell import PDG2SeqCell import numpy as np class PDG2Seq_Encoder(nn.Module): def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim, time_dim, num_layers=1): super(PDG2Seq_Encoder, self).__init__() assert num_layers >= 1, 'At least one DCRNN layer in the Encoder.' self.node_num = node_num self.input_dim = dim_in self.num_layers = num_layers self.PDG2Seq_cells = nn.ModuleList() self.PDG2Seq_cells.append(PDG2SeqCell(node_num, dim_in, dim_out, cheb_k, embed_dim, time_dim)) for _ in range(1, num_layers): self.PDG2Seq_cells.append(PDG2SeqCell(node_num, dim_out, dim_out, cheb_k, embed_dim, time_dim)) def forward(self, x, init_state, node_embeddings): #shape of x: (B, T, N, D) #shape of init_state: (num_layers, B, N, hidden_dim) assert x.shape[2] == self.node_num and x.shape[3] == self.input_dim seq_length = x.shape[1] #x=[batch,steps,nodes,input_dim] current_inputs = x output_hidden = [] for i in range(self.num_layers): state = init_state[i] #state=[batch,steps,nodes,input_dim] inner_states = [] for t in range(seq_length): #如果有两层GRU，则第二层的GGRU的输入是前一层的隐藏状态 state = self.PDG2Seq_cells[i](current_inputs[:, t, :, :], state, [node_embeddings[0][:, t, :], node_embeddings[1][:, t, :], node_embeddings[2]])#state=[batch,steps,nodes,input_dim] # state = self.dcrnn_cells[i](current_inputs[:, t, :, :], state,[node_embeddings[0], node_embeddings[1]]) inner_states.append(state) #一个list，里面是每一步的GRU的hidden状态 output_hidden.append(state) #每层最后一个GRU单元的hidden状态 current_inputs = torch.stack(inner_states, dim=1) #拼接成完整的上一层GRU的hidden状态，作为下一层GRRU的输入[batch,steps,nodes,hiddensize] #current_inputs: the outputs of last layer: (B, T, N, hidden_dim) #output_hidden: the last state for each layer: (num_layers, B, N, hidden_dim) #last_state: (B, N, hidden_dim) return current_inputs, output_hidden def init_hidden(self, batch_size): init_states = [] for i in range(self.num_layers): init_states.append(self.PDG2Seq_cells[i].init_hidden_state(batch_size)) return torch.stack(init_states, dim=0) #(num_layers, B, N, hidden_dim) class PDG2Seq_Dncoder(nn.Module): def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim, time_dim, num_layers=1): super(PDG2Seq_Dncoder, self).__init__() assert num_layers >= 1, 'At least one DCRNN layer in the Decoder.' self.node_num = node_num self.input_dim = dim_in self.num_layers = num_layers self.PDG2Seq_cells = nn.ModuleList() self.PDG2Seq_cells.append(PDG2SeqCell(node_num, dim_in, dim_out, cheb_k, embed_dim, time_dim)) for _ in range(1, num_layers): self.PDG2Seq_cells.append(PDG2SeqCell(node_num, dim_in, dim_out, cheb_k, embed_dim, time_dim)) def forward(self, xt, init_state, node_embeddings): # xt: (B, N, D) # init_state: (num_layers, B, N, hidden_dim) assert xt.shape[1] == self.node_num and xt.shape[2] == self.input_dim current_inputs = xt output_hidden = [] for i in range(self.num_layers): state = self.PDG2Seq_cells[i](current_inputs, init_state[i], [node_embeddings[0], node_embeddings[1], node_embeddings[2]]) output_hidden.append(state) current_inputs = state return current_inputs, output_hidden class PDG2Seq(nn.Module): def __init__(self, args): super(PDG2Seq, self).__init__() self.num_node = args['num_nodes'] self.input_dim = args['input_dim'] self.hidden_dim = args['rnn_units'] self.output_dim = args['output_dim'] self.horizon = args['horizon'] self.num_layers = args['num_layers'] self.use_D = args['use_day'] self.use_W = args['use_week'] self.cl_decay_steps = args['lr_decay_step'] self.node_embeddings1 = nn.Parameter(torch.empty(self.num_node, args['embed_dim'])) self.T_i_D_emb1 = nn.Parameter(torch.empty(288, args['time_dim'])) self.D_i_W_emb1 = nn.Parameter(torch.empty(7, args['time_dim'])) self.T_i_D_emb2 = nn.Parameter(torch.empty(288, args['time_dim'])) self.D_i_W_emb2 = nn.Parameter(torch.empty(7, args['time_dim'])) self.encoder = PDG2Seq_Encoder(args['num_nodes'], args['input_dim'], args['rnn_units'], args['cheb_k'], args['embed_dim'], args['time_dim'], args['num_layers']) self.decoder = PDG2Seq_Dncoder(args['num_nodes'], args['input_dim'], args['rnn_units'], args['cheb_k'], args['embed_dim'], args['time_dim'], args['num_layers']) self.proj = nn.Sequential(nn.Linear(self.hidden_dim, self.output_dim, bias=True)) self.end_conv = nn.Conv2d(1, args['horizon'] * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True) def forward(self, source, traget=None, batches_seen=None): #source: B, T_1, N, D #target: B, T_2, N, D t_i_d_data1 = source[..., 0,-2] t_i_d_data2 = traget[..., 0,-2] # T_i_D_emb = self.T_i_D_emb[(t_i_d_data[:, -1, :] * 288).type(torch.LongTensor)] T_i_D_emb1_en = self.T_i_D_emb1[(t_i_d_data1 * 288).type(torch.LongTensor)] T_i_D_emb2_en = self.T_i_D_emb2[(t_i_d_data1 * 288).type(torch.LongTensor)] T_i_D_emb1_de = self.T_i_D_emb1[(t_i_d_data2 * 288).type(torch.LongTensor)] T_i_D_emb2_de = self.T_i_D_emb2[(t_i_d_data2 * 288).type(torch.LongTensor)] if self.use_W: d_i_w_data1 = source[..., 0,-1] d_i_w_data2 = traget[..., 0,-1] # D_i_W_emb = self.D_i_W_emb[(d_i_w_data[:, -1, :]).type(torch.LongTensor)] D_i_W_emb1_en = self.D_i_W_emb1[(d_i_w_data1).type(torch.LongTensor)] D_i_W_emb2_en = self.D_i_W_emb2[(d_i_w_data1).type(torch.LongTensor)] D_i_W_emb1_de = self.D_i_W_emb1[(d_i_w_data2).type(torch.LongTensor)] D_i_W_emb2_de = self.D_i_W_emb2[(d_i_w_data2).type(torch.LongTensor)] node_embedding_en1 = torch.mul(T_i_D_emb1_en, D_i_W_emb1_en) node_embedding_en2 = torch.mul(T_i_D_emb2_en, D_i_W_emb2_en) node_embedding_de1 = torch.mul(T_i_D_emb1_de, D_i_W_emb1_de) node_embedding_de2 = torch.mul(T_i_D_emb2_de, D_i_W_emb2_de) else: node_embedding_en1 = T_i_D_emb1_en node_embedding_en2 = T_i_D_emb2_en node_embedding_de1 = T_i_D_emb1_de node_embedding_de2 = T_i_D_emb2_de en_node_embeddings=[node_embedding_en1, node_embedding_en2, self.node_embeddings1] source = source[..., 0].unsqueeze(-1) init_state = self.encoder.init_hidden(source.shape[0]).to(source.device) # [2,64,307,64] 前面是2是因为有两层GRU state, _ = self.encoder(source, init_state, en_node_embeddings) # B, T, N, hidden state = state[:, -1:, :, :].squeeze(1) ht_list = [state] * self.num_layers go = torch.zeros((source.shape[0], self.num_node, self.output_dim), device=source.device) out = [] for t in range(self.horizon): state, ht_list = self.decoder(go, ht_list, [node_embedding_de1[:, t, :], node_embedding_de2[:, t, :], self.node_embeddings1]) go = self.proj(state) out.append(go) if self.training: #这里的课程学习用了给予一定概率用真实值代替预测值来学习的教师-学生学习法（名字忘了，大概跟着有关） c = np.random.uniform(0, 1) if c < self._compute_sampling_threshold(batches_seen): #如果满足条件，则用真实值代替预测值训练 go = traget[:, t, :, 0].unsqueeze(-1) output = torch.stack(out, dim=1) return output def _compute_sampling_threshold(self, batches_seen): x = self.cl_decay_steps / ( self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps)) return x