如何实现逐Batch元素采用不同权重的无循环Conv1d操作？

给定输入张量x形状为(32, 4, 16)（对应batch size=32、输入通道数=4、通道维度=16），需要为每个batch元素分配独立的卷积核weight（形状(32, 8, 4, 1)），实现无for循环的Conv1d操作。

方法一：纯PyTorch实现（维度变换+批量矩阵乘法） #

由于卷积核尺寸为1，Conv1d的本质是对每个位置的输入通道做加权求和，可通过torch.bmm（批量矩阵乘法）直接实现：

import torch
import torch.nn.functional as F

# 初始化参数与张量
batch_size = 32
input_channels = 4
output_channels = 8
dim = 16

x = torch.zeros(size=(batch_size, input_channels, dim))
weight = torch.rand(size=(batch_size, output_channels, input_channels, 1))

# 调整卷积核形状：移除最后一维的kernel_size=1维度
weight_reshaped = weight.squeeze(-1)  # 形状变为 (32, 8, 4)
# 执行批量矩阵乘法，得到结果
y = torch.bmm(weight_reshaped, x)  # 输出形状为 (32, 8, 16)

# 验证结果正确性（与循环实现对比）
y_loop = torch.zeros(batch_size, output_channels, dim)
for i in range(batch_size):
    y_loop[i] = F.conv1d(x[i:i+1], weight[i])
print(torch.allclose(y, y_loop))  # 输出True，说明结果一致

方法二：使用einops简化维度操作 #

借助einops库可以更直观地处理维度变换，代码可读性更强（需先安装库：pip install einops）：

from einops import rearrange

# 用einops重排维度后执行矩阵乘法
y_einops = rearrange(weight, 'b o i 1 -> b o i') @ rearrange(x, 'b i d -> b i d')

# 验证结果
print(torch.allclose(y_einops, y_loop))  # 输出True

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