如何在PyTorch中构建支持1xD/NxD核垂直滑动的MxD矩阵输入CNN

实现思路 #

你需要的卷积本质是固定宽度匹配输入列数、仅在行方向滑动的2D卷积，适配2D输入的核心是对输入做维度扩展匹配PyTorch的Conv2d接口要求，具体实现要点如下：

• 输入适配：PyTorch内置Conv2d默认要求4D输入格式(批量大小, 通道数, 高度, 宽度)，因此需要把输入的M×D 2D矩阵扩展两个维度，变为(1, 1, M, D)格式
• 卷积参数配置：卷积核尺寸设置为(N, D)，宽度固定和输入列数一致，仅在行方向做滑动；通过调整N的取值即可切换卷积核高度
• 输出适配：卷积完成后压缩多余的通道、宽度维度，即可得到M×1的目标输出

完整实现代码 #

import torch
import torch.nn as nn

class RowWiseConvCNN(nn.Module):
    def __init__(self, input_col_dim, kernel_row_size=1, use_bias=True):
        """
        input_col_dim: 输入矩阵的列数D
        kernel_row_size: 卷积核的行数N，默认值为1
        use_bias: 是否启用卷积偏置项
        """
        super().__init__()
        self.conv_layer = nn.Conv2d(
            in_channels=1,
            out_channels=1,
            kernel_size=(kernel_row_size, input_col_dim),
            stride=(1, 1),
            # 卷积高度方向padding设置为kernel_row_size//2，保证输出行数和输入一致
            padding=(kernel_row_size//2, 0),
            bias=use_bias
        )
    
    def forward(self, input_mat):
        # input_mat为形状(M, D)的2D输入矩阵
        # 扩展维度适配Conv2d输入要求：(batch=1, channel=1, M, D)
        x = input_mat.unsqueeze(0).unsqueeze(0)
        conv_res = self.conv_layer(x)
        # 压缩多余维度，输出形状为(M, 1)
        return conv_res.squeeze()[:, None]

使用示例 #

# 测试参数：输入矩阵15行8列(M=15, D=8)
M = 15
D = 8
test_input = torch.randn(M, D)

# 1、使用1×D卷积核
model_n1 = RowWiseConvCNN(input_col_dim=D, kernel_row_size=1)
output_n1 = model_n1(test_input)
print(f"N=1时输出形状：{output_n1.shape}") # 输出torch.Size([15, 1])

# 2、调整为3×D卷积核
model_n3 = RowWiseConvCNN(input_col_dim=D, kernel_row_size=3)
output_n3 = model_n3(test_input)
print(f"N=3时输出形状：{output_n3.shape}") # 输出torch.Size([15, 1])

补充说明 #

• 如果需要批量处理多个输入矩阵，可调整前向传播逻辑，输入改为(batch_size, M, D)格式，对应维度扩展时仅增加通道维度即可，最终输出会变为(batch_size, M, 1)
• 如果不需要输出行数和输入保持一致，将Conv2d的padding参数设置为(0, 0)即可，此时输出行数为 M - N + 1

内容的提问来源于stack exchange，提问作者user12314098