class Model_Adp(nn.Module): def __init__(self, SE, args, window_size = 3, T = 12, N=None): super(Model_Adp, self).__init__() L = args.L K = args.K d = args.d D = K * d self.num_his = args.num_his self.SE = SE.to(device) emb_dim = SE.shape[1] self.STEmbedding = STEmbedding(D, emb_dim=emb_dim).to(device) self.STAttBlock_1 = nn.ModuleList([ST_Layer(K, d, T=T, window_size = window_size,N=N) for _ in range(L)]) self.STAttBlock_2 = nn.ModuleList([ST_Layer(K, d, T=T, window_size = window_size,N=N) for _ in range(L)]) self.transformAttention = TransformAttention(K, d) self.mlp_1 = CONVs(input_dims=[1, D], units=[D, D], activations=[F.relu, None]) self.mlp_2 = CONVs(input_dims=[D, D], units=[D, 1], activations=[F.relu, None]) def forward(self, X, TE): # input X = torch.unsqueeze(X, -1) X = self.mlp_1(X) # STE STE = self.STEmbedding(self.SE, TE) STE_his = STE[:, :self.num_his] STE_pred = STE[:, self.num_his:] # encoder for net in self.STAttBlock_1: X = net(X, STE_his) X = self.transformAttention(X, STE_his, STE_pred) # decoder for net in self.STAttBlock_2: X = net(X, STE_pred) X = self.mlp_2(X) del STE, STE_his, STE_pred return torch.squeeze(X, 3)
时间: 2023-06-02 19:06:18 浏览: 138
MPI_hello_world.zip_use mpi_world
这段代码定义了一个名为Model_Adp的类,继承自nn.Module。在初始化函数中,需要传入SE、args、window_size、T和N等参数。该类包含多个属性和方法,其中包括一个STEmbedding、两个STAttBlock、一个TransformAttention、两个CONVs和一个forward方法。其中STEmbedding是对空间数据做embedding操作,STAttBlock是空间-时间数据的多头注意力机制,TransformAttention是将多头注意力输出转化为特征向量,CONVs是一些卷积神经网络。forward方法用来执行前向传播。具体实现可以参考代码中的实现。
阅读全文