关于torch.nn.Linear的输入与输出探讨

关于torch.nn.Linear的一点问题

  最近写Transformer代码的时候，FFN层是线性层，突然发现我丢进去的是一个三维的张量[batch_size, H, d_model]，但是线性层nn.Linear竟然还给它算出来了，而且预期结果形状还是我想要的[batch_size, H, d_ff]。印象中nn.Linear的输入和输出都是二维的。于是我去找了一圈博客，各种博客里面写的都是输入二维的，更疑惑了。最后去看了官方文档才知道，他们都有问题。

  输入和输出的维度可以是任意，而且输出的形状除了最后一个维度，其他的均与输出一样。比如[32, 128, 192]形状的张量，通过nn.Linear(192, 256)的线性层，其输出的形状是[32, 128, 256]。