class Encoder(nn.Module): def __init__(self): super(Encoder, self).__init__() self.fc1 = nn.Sequential( nn.Linear(200, 128), nn.BatchNorm1d(num_features=128), nn.ReLU() ) self.lstm = nn.LSTM(input_size=128, hidden_size=256, num_layers=2, batch_first=True) # (S,N,V) -> (N,S,V) def forward(self, x): # (N,3,42,130) -> (N,126,130) -> (N,130,126) -> (N*130,126) -> (N*130,128) -> (N,130,128) -> (N,128) -> (N,256) # (N,3,50,200) -> (N,150,200) -> (N,200,150) -> (N*200,150) -> (N*200,128) -> (N,200,128) -> (N,128) -> (N,256) x = x.reshape(-1, 150, 200).permute(0, 2, 1) x = x.reshape(-1, 150) fc1 = self.fc1(x) fc1 = fc1.reshape(-1, 200, 128) lstm, (h_n, h_c) = self.lstm(fc1, None) out = lstm[:, -1, :] return out class Decoder(nn.Module): def __init__(self): super(Decoder, self).__init__() self.lstm = nn.LSTM(input_size=256, hidden_size=128, num_layers=2, batch_first=True) self.out = nn.Linear(128, 36) def forward(self, x): # (N,256) -> (N,7,256) -> (N,7,128) -> (N*7,128) -> (N*7,36) -> (N,7,36) # (N,256) -> (N,7,256) -> (N,7,128) -> (N*7,128) -> (N*7,36) -> (N,7,36) x = x.reshape(-1, 1, 256) x = x.expand(-1, 7, 256) lstm, (h_n, h_c) = self.lstm(x, None) y1 = lstm.reshape(-1, 128) out = self.out(y1) output = out.reshape(-1, 7, 36) return output class MainNet (nn.Module): def __init__(self): super(MainNet, self).__init__() self.encoder = Encoder() self.decoder = Decoder() def forward(self, x): encoder = self.encoder(x) decoder = self.decoder(encoder) return decoder

这段代码定义了一个由三个子模块组成的神经网络，其中包括一个编码器、一个解码器和一个主网络。编码器将形状为 (N, 3, 50, 200) 的输入张量 x 通过一系列操作转换为形状为 (N, 256) 的隐状态向量，解码器则将该向量作为输入，通过 LSTM 层和全连接层将其转换为形状为 (N, 7, 36) 的输出张量。主网络将编码器和解码器组合起来，将输入张量传递给编码器，然后将编码器的输出传递给解码器，最终返回解码器的输出。

具体来说，编码器由一个全连接层、一个批归一化层和一个 ReLU 激活函数组成，用于将输入张量进行特征提取。然后，将形状为 (N, 150, 200) 的张量重塑为形状为 (N, 200, 150) 的张量，然后交换第二个和第三个维度，将形状变为 (N, 130, 128)。接着，将张量再次重塑为形状为 (N*130, 128) 的二维张量，并通过全连接层将其转换为形状为 (N, 130, 128) 的三维张量。最后，将张量沿第一个维度输入到 LSTM 层中，将其转换为形状为 (N, 256) 的隐状态向量。

解码器由一个 LSTM 层和一个全连接层组成，用于将编码器的输出转换为形状为 (N, 7, 36) 的输出张量。具体来说，将形状为 (N, 256) 的隐状态向量重复 7 次，然后通过 LSTM 层将其转换为形状为 (N, 7, 128) 的三维张量。接着，将张量重塑为形状为 (N*7, 128) 的二维张量，并通过全连接层将其转换为形状为 (N*7, 36) 的二维张量。最后，将张量重塑为形状为 (N, 7, 36) 的三维张量作为输出。

主网络由一个编码器和一个解码器组成，用于将输入张量传递给编码器，然后将编码器的输出传递给解码器，最终返回解码器的输出。具体来说，将输入张量传递给编码器的 forward 函数，得到编码器的输出。然后，将编码器的输出传递给解码器的 forward 函数，得到解码器的输出，并将其返回。