neural-network - 如何仅在一个维度(按行或按列)将线性层应用于二维层 - 部分连接层

Question

我正在尝试将线性层应用于仅按列连接的张量的二维矩阵，如下图所示。

输入形状是(batch_size, 3, 50)。我首先尝试使用 2D 卷积，添加 1 channel 维度，因此输入形状为 (batch_size, 1, 3, 50)

import torch.nn as nn
import torch

class ColumnConv(nn.Module):
    def __init__(self):
        self.layers = nn.Sequential(
            nn.Conv2d(
                in_channels=1,
                out_channels=1,
                kernel_size=(3, 1),
                stride=1,
            ),  # shape is B, 1, 1, 50
            nn.ReLU(),
            nn.Flatten() #shape is B, 50
        )

    def forward(self, x):
        return self.layers(x)

但是好像不行。我计划使用 50 个 nn.Linear 层的列表并将它们应用于输入的列切片，但它看起来更像是一种未针对性能进行优化的解决方法。

是否有更“pytorchic”的方式来做到这一点？

Answer 1

PyTorch nn.Linear模块可以应用于多维输入，线性将应用于最后一个维度，因此按列应用解决方案是交换行和列。

linear_3_to_1 = nn.Linear(3, 1)

x = torch.randn(1, 1, 3, 50)
x = x.transpose(2, 3) #swap 3 and 50
out = linear_3_to_1(x).flatten()