如何在PyTorch中正确定义权重与偏置以获得正确形状输出？

错误根源 #

1. 冗余循环无意义且引发维度问题：原forward方法中的循环完全多余，每次循环都处理整个输入张量，最终仅保留最后一次循环结果，且循环内的展平操作会混淆批量维度与特征维度。
2. 输入张量形状处理错误：使用flatten会将整个张量压成一维，丢失批量维度，而线性层要求输入为二维（批量大小×特征数）。
3. 最后一层权重形状不兼容：nn.Linear(3,1)的权重形状应为[1,3]，但传入的weight_last是[3]，导致矩阵乘法维度不匹配。

修正后的代码 #

模型初始化代码 #

# 定义神经网络
def __init__(self, weight, bias, weight_last, bias_last):
    # weight.shape = [3,3,3]
    # bias.shape = [3,3]
    # weight_last 需要调整为 [1,3]，对应最后一层nn.Linear(3,1)的权重形状
    # bias_last = [1]，符合最后一层偏置形状
    
    super(NeuralNetHardeningModel, self).__init__()
    
    self.weight = weight
    self.bias = bias
    self.weight_last = weight_last
    self.bias_last = bias_last
    
    self.nn = nn.Sequential(
        nn.Linear(3, 3),
        nn.ReLU(),
        nn.Linear(3, 3),
        nn.ReLU(),
        nn.Linear(3, 3),
        nn.ReLU(),
        nn.Linear(3, 1)
    )
    
    if len(weight.shape) == 3:
        with torch.no_grad():
            self.nn[0].weight = nn.Parameter(weight[0])
            self.nn[0].bias = nn.Parameter(bias[0])
            
            self.nn[2].weight = nn.Parameter(weight[1])
            self.nn[2].bias = nn.Parameter(bias[1])
            
            self.nn[4].weight = nn.Parameter(weight[2])
            self.nn[4].bias = nn.Parameter(bias[2])
            
            # 自动调整weight_last形状为[1,3]（若原先是[3]）
            self.nn[6].weight = nn.Parameter(weight_last.unsqueeze(0) if weight_last.dim() == 1 else weight_last)
            self.nn[6].bias = nn.Parameter(bias_last)

前向传播代码 #

# 神经网络前向传播方法
def forward(self, a, b, c):
    # 去掉输入张量中多余的单维度，保留[批量大小, 1]的形状
    a_eval = a.squeeze(-1) if a.dim() > 2 else a
    b_eval = b.squeeze(-1) if b.dim() > 2 else b
    c_eval = c.squeeze(-1) if c.dim() > 2 else c
    
    # 在特征维度拼接，得到[70,3]的批量输入
    y = torch.cat((a_eval, b_eval, c_eval), dim=1)
    
    # 直接传入网络，得到[70,1]的输出
    y1 = self.nn(y)
    
    return y1

关键说明 #

• 移除冗余循环：PyTorch线性层原生支持批量输入，只要输入为[N, in_features]格式，即可直接处理并输出[N, out_features]。
• 形状调整：用squeeze(-1)替代flatten，仅去除多余的单维度，保留批量维度与特征维度的区分。
• 权重适配：确保最后一层权重形状符合nn.Linear(out_features, in_features)的要求，通过unsqueeze(0)自动修正一维权重的形状。

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